267 lines
8.4 KiB
Python
267 lines
8.4 KiB
Python
"""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 paths below, or pass them on the command line:
|
|
|
|
python match_stamp.py --image output\20260623_094529\TY037\1_TY037.png --stamp-dir D:\stamps
|
|
|
|
It does not copy, move, rename, delete, or write image files. Results are
|
|
printed directly.
|
|
"""
|
|
from __future__ import print_function
|
|
|
|
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
|
|
|
|
|
|
MERGED_IMAGE = r"D:\chengma\cmbot\output\20260623_094529\TY037\1_TY037.png"
|
|
STAMP_DIR = r"D:\chengma\印花和底图\已处理印花\卡通71（66大码200斤 KEKE已上）\横1"
|
|
|
|
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
|
|
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 iter_images(root):
|
|
root = Path(root)
|
|
if not root.exists():
|
|
return
|
|
for path in sorted(root.rglob("*")):
|
|
if is_image(path):
|
|
yield path
|
|
|
|
|
|
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
|
|
|
|
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 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()
|
|
|
|
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)
|
|
|
|
results.sort(key=lambda item: item["score"], reverse=True)
|
|
return results[:top_n], len(results)
|
|
|
|
|
|
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 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 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
|
|
|
|
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__":
|
|
raise SystemExit(main())
|