feat: support nonuniform stamp matching
This commit is contained in:
+50
-16
@@ -34,6 +34,7 @@ TOP_N = 5
|
||||
SCALE_MIN = 0.10
|
||||
SCALE_MAX = 2.00
|
||||
SCALE_STEP = 0.05
|
||||
ASPECT_Y_FACTORS = (0.60, 0.70, 0.80, 0.90, 1.00, 1.10, 1.20)
|
||||
|
||||
ALPHA_THRESHOLD = 20
|
||||
LOW_SCORE_WARNING_THRESHOLD = 0.55
|
||||
@@ -96,6 +97,27 @@ def build_scales():
|
||||
return scales
|
||||
|
||||
|
||||
def build_scale_variants():
|
||||
variants = []
|
||||
seen = set()
|
||||
for scale in build_scales():
|
||||
for aspect_y in ASPECT_Y_FACTORS:
|
||||
scale_x = scale
|
||||
scale_y = round(scale * aspect_y, 4)
|
||||
key = (scale_x, scale_y, aspect_y)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
variants.append(
|
||||
{
|
||||
"scale_x": scale_x,
|
||||
"scale_y": scale_y,
|
||||
"aspect_y": aspect_y,
|
||||
}
|
||||
)
|
||||
return variants
|
||||
|
||||
|
||||
def safe_match_template(source, template, mask=None):
|
||||
if template.shape[0] > source.shape[0] or template.shape[1] > source.shape[1]:
|
||||
return None
|
||||
@@ -119,25 +141,28 @@ def safe_match_template(source, template, mask=None):
|
||||
return np.clip(result, 0.0, 1.0)
|
||||
|
||||
|
||||
def resize_for_scale(image, scale, interpolation):
|
||||
def resize_for_scale(image, scale_x, scale_y, interpolation):
|
||||
height, width = image.shape[:2]
|
||||
new_width = max(1, int(round(width * scale)))
|
||||
new_height = max(1, int(round(height * scale)))
|
||||
new_width = max(1, int(round(width * scale_x)))
|
||||
new_height = max(1, int(round(height * scale_y)))
|
||||
return cv2.resize(image, (new_width, new_height), interpolation=interpolation)
|
||||
|
||||
|
||||
def score_stamp(merged_rgb, merged_edges, stamp_path, scales):
|
||||
def score_stamp(merged_rgb, merged_edges, stamp_path, scale_variants):
|
||||
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)
|
||||
for variant in scale_variants:
|
||||
scale_x = variant["scale_x"]
|
||||
scale_y = variant["scale_y"]
|
||||
aspect_y = variant["aspect_y"]
|
||||
scaled_rgb = resize_for_scale(stamp_rgb, scale_x, scale_y, cv2.INTER_AREA)
|
||||
scaled_gray = resize_for_scale(stamp_gray, scale_x, scale_y, cv2.INTER_AREA)
|
||||
scaled_edges = resize_for_scale(stamp_edges, scale_x, scale_y, cv2.INTER_NEAREST)
|
||||
scaled_mask = resize_for_scale(stamp_mask, scale_x, scale_y, cv2.INTER_NEAREST)
|
||||
|
||||
template_result = safe_match_template(merged_rgb, scaled_rgb, scaled_mask)
|
||||
if template_result is None:
|
||||
@@ -165,7 +190,9 @@ def score_stamp(merged_rgb, merged_edges, stamp_path, scales):
|
||||
"stamp_path": stamp_path,
|
||||
"x": int(x),
|
||||
"y": int(y),
|
||||
"scale": float(scale),
|
||||
"scale_x": float(scale_x),
|
||||
"scale_y": float(scale_y),
|
||||
"aspect_y": float(aspect_y),
|
||||
"width": int(scaled_gray.shape[1]),
|
||||
"height": int(scaled_gray.shape[0]),
|
||||
}
|
||||
@@ -178,12 +205,12 @@ def find_best_matches(merged_image, stamp_dir, top_n):
|
||||
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()
|
||||
scale_variants = build_scale_variants()
|
||||
|
||||
results = []
|
||||
for stamp_path in iter_images(stamp_dir):
|
||||
try:
|
||||
result = score_stamp(merged_rgb, merged_edges, stamp_path, scales)
|
||||
result = score_stamp(merged_rgb, merged_edges, stamp_path, scale_variants)
|
||||
except Exception as exc:
|
||||
print("Skipped: {} ({})".format(stamp_path, exc))
|
||||
continue
|
||||
@@ -211,7 +238,9 @@ def print_result(results, matched_count, merged_image, stamp_dir):
|
||||
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("scale_x: {:.4f}".format(best["scale_x"]))
|
||||
print("scale_y: {:.4f}".format(best["scale_y"]))
|
||||
print("aspect_y: {:.4f}".format(best["aspect_y"]))
|
||||
print("size: {}x{}".format(best["width"], best["height"]))
|
||||
if best["score"] < LOW_SCORE_WARNING_THRESHOLD:
|
||||
print("warning: best score is low; please review manually.")
|
||||
@@ -219,14 +248,19 @@ def print_result(results, matched_count, merged_image, stamp_dir):
|
||||
print("")
|
||||
print("Top {}:".format(len(results)))
|
||||
for index, item in enumerate(results, 1):
|
||||
line = (
|
||||
"{}. {:.4f} template={:.4f} edge={:.4f} "
|
||||
"scale_x={:.4f} scale_y={:.4f} aspect_y={:.4f} {}"
|
||||
)
|
||||
print(
|
||||
"{}. {:.4f} template={:.4f} edge={:.4f} scale={:.4f} {}"
|
||||
.format(
|
||||
line.format(
|
||||
index,
|
||||
item["score"],
|
||||
item["template_score"],
|
||||
item["edge_score"],
|
||||
item["scale"],
|
||||
item["scale_x"],
|
||||
item["scale_y"],
|
||||
item["aspect_y"],
|
||||
Path(item["stamp_path"]).name,
|
||||
)
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user