feat: implement stamp content matching
This commit is contained in:
+222
-112
@@ -1,50 +1,52 @@
|
||||
"""Generate a CSV report matching merged images to source stamp images.
|
||||
"""Find the source stamp image that best matches a merged clothing image.
|
||||
|
||||
This script is intentionally standalone and does not import cmbot project code.
|
||||
Edit the three paths below, then run:
|
||||
Edit the paths below, or pass them on the command line:
|
||||
|
||||
python match_stamp.py
|
||||
python match_stamp.py --image output\20260623_094529\TY037\1_TY037.png --stamp-dir D:\stamps
|
||||
|
||||
It does not copy, move, rename, or delete any image files.
|
||||
It does not copy, move, rename, delete, or write image files. Results are
|
||||
printed directly.
|
||||
"""
|
||||
from __future__ import print_function
|
||||
|
||||
import csv
|
||||
from datetime import datetime
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
except ImportError as exc:
|
||||
print("Missing dependency: {}".format(exc))
|
||||
print("Install required packages: pip install opencv-python numpy Pillow")
|
||||
raise
|
||||
|
||||
OUTPUT_DIR = r"D:\chengma\cmbot\output\20260623_094529"
|
||||
|
||||
MERGED_IMAGE = r"D:\chengma\cmbot\output\20260623_094529\TY037\1_TY037.png"
|
||||
STAMP_DIR = r"D:\chengma\印花和底图\已处理印花\卡通71（66大码200斤 KEKE已上）\横1"
|
||||
REPORT_PATH = r"D:\chengma\cmbot\match_stamp_report.csv"
|
||||
|
||||
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
|
||||
FIELDNAMES = [
|
||||
"status",
|
||||
"code",
|
||||
"merged_image",
|
||||
"stamp_image",
|
||||
"merged_folder",
|
||||
"message",
|
||||
]
|
||||
TOP_N = 5
|
||||
|
||||
SCALE_MIN = 0.10
|
||||
SCALE_MAX = 2.00
|
||||
SCALE_STEP = 0.05
|
||||
|
||||
ALPHA_THRESHOLD = 20
|
||||
LOW_SCORE_WARNING_THRESHOLD = 0.55
|
||||
TEMPLATE_SCORE_WEIGHT = 0.75
|
||||
EDGE_SCORE_WEIGHT = 0.25
|
||||
MIN_MATCH_WIDTH = 24
|
||||
MIN_MATCH_HEIGHT = 24
|
||||
|
||||
|
||||
def is_image(path):
|
||||
return path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
|
||||
|
||||
|
||||
def extract_code_from_merged_name(path):
|
||||
"""Extract stamp code from the final underscore suffix.
|
||||
|
||||
Example: 1_TY037.png -> TY037.
|
||||
"""
|
||||
stem = path.stem
|
||||
if "_" not in stem:
|
||||
return ""
|
||||
code = stem.rsplit("_", 1)[1].strip()
|
||||
return code
|
||||
|
||||
|
||||
def iter_images(root):
|
||||
root = Path(root)
|
||||
if not root.exists():
|
||||
@@ -54,103 +56,211 @@ def iter_images(root):
|
||||
yield path
|
||||
|
||||
|
||||
def build_stamp_index(stamp_dir):
|
||||
index = {}
|
||||
for path in iter_images(stamp_dir):
|
||||
code = path.stem.strip()
|
||||
if not code:
|
||||
def load_rgba(path):
|
||||
"""Load with Pillow so Windows Chinese paths are handled reliably."""
|
||||
with Image.open(str(path)) as image:
|
||||
return np.array(image.convert("RGBA"))
|
||||
|
||||
|
||||
def rgba_to_gray(rgba):
|
||||
rgb = rgba[:, :, :3]
|
||||
return cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
|
||||
|
||||
|
||||
def rgba_to_rgb(rgba):
|
||||
return rgba[:, :, :3]
|
||||
|
||||
|
||||
def crop_transparent_border(rgba):
|
||||
alpha = rgba[:, :, 3]
|
||||
ys, xs = np.where(alpha > ALPHA_THRESHOLD)
|
||||
if len(xs) == 0 or len(ys) == 0:
|
||||
mask = np.ones(alpha.shape, dtype=np.uint8) * 255
|
||||
return rgba[:, :, :3], mask
|
||||
|
||||
left = int(xs.min())
|
||||
right = int(xs.max()) + 1
|
||||
top = int(ys.min())
|
||||
bottom = int(ys.max()) + 1
|
||||
cropped = rgba[top:bottom, left:right]
|
||||
mask = (cropped[:, :, 3] > ALPHA_THRESHOLD).astype(np.uint8) * 255
|
||||
return cropped[:, :, :3], mask
|
||||
|
||||
|
||||
def build_scales():
|
||||
scales = []
|
||||
value = SCALE_MIN
|
||||
while value <= SCALE_MAX + 0.0001:
|
||||
scales.append(round(value, 4))
|
||||
value += SCALE_STEP
|
||||
return scales
|
||||
|
||||
|
||||
def safe_match_template(source, template, mask=None):
|
||||
if template.shape[0] > source.shape[0] or template.shape[1] > source.shape[1]:
|
||||
return None
|
||||
if template.shape[0] < MIN_MATCH_HEIGHT or template.shape[1] < MIN_MATCH_WIDTH:
|
||||
return None
|
||||
if mask is not None and int(np.count_nonzero(mask)) < 9:
|
||||
return None
|
||||
|
||||
match_mask = mask
|
||||
if match_mask is not None and template.ndim == 3 and match_mask.ndim == 2:
|
||||
channels = template.shape[2]
|
||||
match_mask = cv2.merge([match_mask] * channels)
|
||||
|
||||
try:
|
||||
result = cv2.matchTemplate(source, template, cv2.TM_CCORR_NORMED, mask=match_mask)
|
||||
except cv2.error:
|
||||
return None
|
||||
|
||||
result = np.asarray(result, dtype=np.float32)
|
||||
result[~np.isfinite(result)] = -1.0
|
||||
return np.clip(result, 0.0, 1.0)
|
||||
|
||||
|
||||
def resize_for_scale(image, scale, interpolation):
|
||||
height, width = image.shape[:2]
|
||||
new_width = max(1, int(round(width * scale)))
|
||||
new_height = max(1, int(round(height * scale)))
|
||||
return cv2.resize(image, (new_width, new_height), interpolation=interpolation)
|
||||
|
||||
|
||||
def score_stamp(merged_rgb, merged_edges, stamp_path, scales):
|
||||
stamp_rgba = load_rgba(stamp_path)
|
||||
stamp_rgb, stamp_mask = crop_transparent_border(stamp_rgba)
|
||||
stamp_gray = cv2.cvtColor(stamp_rgb, cv2.COLOR_RGB2GRAY)
|
||||
stamp_edges = cv2.Canny(stamp_gray, 80, 160)
|
||||
|
||||
best = None
|
||||
for scale in scales:
|
||||
scaled_rgb = resize_for_scale(stamp_rgb, scale, cv2.INTER_AREA)
|
||||
scaled_gray = resize_for_scale(stamp_gray, scale, cv2.INTER_AREA)
|
||||
scaled_edges = resize_for_scale(stamp_edges, scale, cv2.INTER_NEAREST)
|
||||
scaled_mask = resize_for_scale(stamp_mask, scale, cv2.INTER_NEAREST)
|
||||
|
||||
template_result = safe_match_template(merged_rgb, scaled_rgb, scaled_mask)
|
||||
if template_result is None:
|
||||
continue
|
||||
index.setdefault(code, []).append(path.resolve())
|
||||
return index
|
||||
|
||||
edge_result = safe_match_template(merged_edges, scaled_edges, scaled_mask)
|
||||
if edge_result is None:
|
||||
edge_result = np.zeros(template_result.shape, dtype=np.float32)
|
||||
|
||||
combined = (
|
||||
TEMPLATE_SCORE_WEIGHT * template_result
|
||||
+ EDGE_SCORE_WEIGHT * edge_result
|
||||
)
|
||||
_min_val, max_val, _min_loc, max_loc = cv2.minMaxLoc(combined)
|
||||
x, y = max_loc
|
||||
template_score = float(template_result[y, x])
|
||||
edge_score = float(edge_result[y, x])
|
||||
score = float(max_val)
|
||||
|
||||
if best is None or score > best["score"]:
|
||||
best = {
|
||||
"score": score,
|
||||
"template_score": template_score,
|
||||
"edge_score": edge_score,
|
||||
"stamp_path": stamp_path,
|
||||
"x": int(x),
|
||||
"y": int(y),
|
||||
"scale": float(scale),
|
||||
"width": int(scaled_gray.shape[1]),
|
||||
"height": int(scaled_gray.shape[0]),
|
||||
}
|
||||
|
||||
return best
|
||||
|
||||
|
||||
def match_merged_image(path, stamp_index):
|
||||
code = extract_code_from_merged_name(path)
|
||||
if not code:
|
||||
return {
|
||||
"status": "bad_name",
|
||||
"code": "",
|
||||
"merged_image": str(path.resolve()),
|
||||
"stamp_image": "",
|
||||
"merged_folder": str(path.parent.resolve()),
|
||||
"message": "File name has no underscore suffix code",
|
||||
}
|
||||
def find_best_matches(merged_image, stamp_dir, top_n):
|
||||
merged_rgba = load_rgba(merged_image)
|
||||
merged_rgb = rgba_to_rgb(merged_rgba)
|
||||
merged_gray = rgba_to_gray(merged_rgba)
|
||||
merged_edges = cv2.Canny(merged_gray, 80, 160)
|
||||
scales = build_scales()
|
||||
|
||||
candidates = stamp_index.get(code, [])
|
||||
if not candidates:
|
||||
return {
|
||||
"status": "missing",
|
||||
"code": code,
|
||||
"merged_image": str(path.resolve()),
|
||||
"stamp_image": "",
|
||||
"merged_folder": str(path.parent.resolve()),
|
||||
"message": "No stamp image found for code",
|
||||
}
|
||||
results = []
|
||||
for stamp_path in iter_images(stamp_dir):
|
||||
try:
|
||||
result = score_stamp(merged_rgb, merged_edges, stamp_path, scales)
|
||||
except Exception as exc:
|
||||
print("Skipped: {} ({})".format(stamp_path, exc))
|
||||
continue
|
||||
if result is not None:
|
||||
results.append(result)
|
||||
|
||||
stamp_paths = " | ".join(str(p) for p in candidates)
|
||||
if len(candidates) > 1:
|
||||
return {
|
||||
"status": "duplicate",
|
||||
"code": code,
|
||||
"merged_image": str(path.resolve()),
|
||||
"stamp_image": stamp_paths,
|
||||
"merged_folder": str(path.parent.resolve()),
|
||||
"message": "Multiple stamp images found for code",
|
||||
}
|
||||
|
||||
return {
|
||||
"status": "matched",
|
||||
"code": code,
|
||||
"merged_image": str(path.resolve()),
|
||||
"stamp_image": stamp_paths,
|
||||
"merged_folder": str(path.parent.resolve()),
|
||||
"message": "",
|
||||
}
|
||||
results.sort(key=lambda item: item["score"], reverse=True)
|
||||
return results[:top_n], len(results)
|
||||
|
||||
|
||||
def write_csv(rows, report_path):
|
||||
report_path = Path(report_path)
|
||||
report_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with report_path.open("w", encoding="utf-8-sig", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=FIELDNAMES)
|
||||
writer.writeheader()
|
||||
for row in rows:
|
||||
writer.writerow(row)
|
||||
def print_result(results, matched_count, merged_image, stamp_dir):
|
||||
print("Merged image: {}".format(Path(merged_image).resolve()))
|
||||
print("Stamp dir: {}".format(Path(stamp_dir).resolve()))
|
||||
print("Matched stamp candidates: {}".format(matched_count))
|
||||
print("")
|
||||
|
||||
if not results:
|
||||
print("No valid image match result.")
|
||||
return
|
||||
|
||||
best = results[0]
|
||||
print("Best match:")
|
||||
print("score: {:.4f}".format(best["score"]))
|
||||
print("template_score: {:.4f}".format(best["template_score"]))
|
||||
print("edge_score: {:.4f}".format(best["edge_score"]))
|
||||
print("stamp: {}".format(Path(best["stamp_path"]).resolve()))
|
||||
print("location: x={}, y={}".format(best["x"], best["y"]))
|
||||
print("scale: {:.4f}".format(best["scale"]))
|
||||
print("size: {}x{}".format(best["width"], best["height"]))
|
||||
if best["score"] < LOW_SCORE_WARNING_THRESHOLD:
|
||||
print("warning: best score is low; please review manually.")
|
||||
|
||||
print("")
|
||||
print("Top {}:".format(len(results)))
|
||||
for index, item in enumerate(results, 1):
|
||||
print(
|
||||
"{}. {:.4f} template={:.4f} edge={:.4f} scale={:.4f} {}"
|
||||
.format(
|
||||
index,
|
||||
item["score"],
|
||||
item["template_score"],
|
||||
item["edge_score"],
|
||||
item["scale"],
|
||||
Path(item["stamp_path"]).name,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def generate_report(output_dir, stamp_dir, report_path):
|
||||
output_dir = Path(output_dir)
|
||||
stamp_dir = Path(stamp_dir)
|
||||
|
||||
stamp_index = build_stamp_index(stamp_dir)
|
||||
rows = []
|
||||
for path in iter_images(output_dir):
|
||||
rows.append(match_merged_image(path, stamp_index))
|
||||
|
||||
write_csv(rows, report_path)
|
||||
return rows
|
||||
def parse_args(argv):
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Match one merged clothing image against a stamp directory."
|
||||
)
|
||||
parser.add_argument("--image", default=MERGED_IMAGE, help="Merged clothing image path")
|
||||
parser.add_argument("--stamp-dir", default=STAMP_DIR, help="Stamp image directory")
|
||||
parser.add_argument("--top", default=TOP_N, type=int, help="Number of results to print")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def summarize(rows):
|
||||
counts = {}
|
||||
for row in rows:
|
||||
status = row.get("status", "")
|
||||
counts[status] = counts.get(status, 0) + 1
|
||||
return counts
|
||||
def main(argv=None):
|
||||
args = parse_args(argv or sys.argv[1:])
|
||||
merged_image = Path(args.image)
|
||||
stamp_dir = Path(args.stamp_dir)
|
||||
|
||||
if not merged_image.is_file():
|
||||
print("Merged image not found: {}".format(merged_image))
|
||||
return 1
|
||||
if not stamp_dir.is_dir():
|
||||
print("Stamp directory not found: {}".format(stamp_dir))
|
||||
return 1
|
||||
|
||||
def main():
|
||||
started = datetime.now()
|
||||
rows = generate_report(OUTPUT_DIR, STAMP_DIR, REPORT_PATH)
|
||||
counts = summarize(rows)
|
||||
print("Generated report: {}".format(Path(REPORT_PATH).resolve()))
|
||||
print("Merged images: {}".format(len(rows)))
|
||||
print("Matched: {}".format(counts.get("matched", 0)))
|
||||
print("Missing: {}".format(counts.get("missing", 0)))
|
||||
print("Duplicate: {}".format(counts.get("duplicate", 0)))
|
||||
print("Bad name: {}".format(counts.get("bad_name", 0)))
|
||||
print("Elapsed: {}".format(datetime.now() - started))
|
||||
started = time.time()
|
||||
results, matched_count = find_best_matches(merged_image, stamp_dir, max(1, args.top))
|
||||
print_result(results, matched_count, merged_image, stamp_dir)
|
||||
print("")
|
||||
print("Elapsed: {:.2f}s".format(time.time() - started))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
raise SystemExit(main())
|
||||
|
||||
Reference in New Issue
Block a user