150 lines
5.1 KiB
Python
150 lines
5.1 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Brute Force Seal Unwarping
|
|
Tests all permutations of rotation and flipping to find 70%+ similarity.
|
|
"""
|
|
import sys
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
import difflib
|
|
|
|
# Target text
|
|
TARGET_TEXT = "威凯检测技术有限公司"
|
|
MIN_SIMILARITY = 0.70
|
|
|
|
def similarity(s1, s2):
|
|
if not s1 or not s2: return 0.0
|
|
return difflib.SequenceMatcher(None, s1, s2).ratio()
|
|
|
|
def brute_force_search(image_path, results_dir):
|
|
if not os.path.exists(results_dir):
|
|
os.makedirs(results_dir)
|
|
|
|
img = cv2.imread(image_path)
|
|
if img is None:
|
|
print(f"Error: Could not load {image_path}")
|
|
return
|
|
|
|
# 1. Circle Detection
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
gray = cv2.medianBlur(gray, 5)
|
|
rows = gray.shape[0]
|
|
circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, 1, rows / 8,
|
|
param1=100, param2=30,
|
|
minRadius=int(rows / 4), maxRadius=int(rows / 1.5))
|
|
|
|
if circles is not None:
|
|
best_circle = np.uint16(np.around(circles))[0][0]
|
|
cx, cy, radius = [int(x) for x in best_circle]
|
|
else:
|
|
h, w = img.shape[:2]
|
|
cx, cy, radius = w // 2, h // 2, min(w, h) // 2
|
|
|
|
print(f"Using Center ({cx}, {cy}), Radius {radius}", file=sys.stderr)
|
|
|
|
# 2. Setup OCR
|
|
try:
|
|
os.environ["DISABLE_MODEL_SOURCE_CHECK"] = "True"
|
|
from paddleocr import TextRecognition
|
|
text_rec = TextRecognition(model_name="PP-OCRv4_server_rec")
|
|
except Exception as e:
|
|
print(f"OCR Error: {e}")
|
|
return
|
|
|
|
# 3. Permutations
|
|
rotations = [
|
|
(None, "no_rot"),
|
|
(cv2.ROTATE_90_CLOCKWISE, "90_cw"),
|
|
(cv2.ROTATE_180, "180"),
|
|
(cv2.ROTATE_90_COUNTERCLOCKWISE, "90_ccw")
|
|
]
|
|
|
|
flips = [
|
|
(None, "no_flip"),
|
|
(0, "flip_v"),
|
|
(1, "flip_h"),
|
|
(-1, "flip_both")
|
|
]
|
|
|
|
# Text ring is usually outer 25% of radius
|
|
factors = [1.5, 2.5] # Test two widths
|
|
|
|
found_best = None
|
|
max_sim = 0.0
|
|
|
|
for rot_code, rot_name in rotations:
|
|
# Rotate image or not
|
|
if rot_code is not None:
|
|
work_img = cv2.rotate(img, rot_code)
|
|
# Re-detect center on rotated image to be exact
|
|
gray_rot = cv2.cvtColor(work_img, cv2.COLOR_BGR2GRAY)
|
|
circles_rot = cv2.HoughCircles(gray_rot, cv2.HOUGH_GRADIENT, 1, rows / 8,
|
|
param1=100, param2=30,
|
|
minRadius=int(radius * 0.8), maxRadius=int(radius * 1.2))
|
|
if circles_rot is not None:
|
|
bc_r = np.uint16(np.around(circles_rot))[0][0]
|
|
rcx, rcy, rr = [int(x) for x in bc_r]
|
|
else:
|
|
rh, rw = work_img.shape[:2]
|
|
rcx, rcy, rr = rw // 2, rh // 2, radius
|
|
else:
|
|
work_img = img.copy()
|
|
rcx, rcy, rr = cx, cy, radius
|
|
|
|
for factor in factors:
|
|
# Warp
|
|
out_w = int(rr * 2 * np.pi * factor)
|
|
out_h = rr
|
|
unwarped = cv2.warpPolar(work_img, (out_w, out_h), (rcx, rcy), rr, cv2.WARP_POLAR_LINEAR)
|
|
|
|
# Crop outer ring
|
|
crop_outer = unwarped[int(out_h * 0.7):out_h, :]
|
|
|
|
for flip_code, flip_name in flips:
|
|
final = crop_outer.copy()
|
|
if flip_code is not None:
|
|
final = cv2.flip(final, flip_code)
|
|
|
|
# Test OCR
|
|
test_name = f"{rot_name}_{flip_name}_f{factor:.1f}"
|
|
temp_path = os.path.join(results_dir, f"{test_name}.png")
|
|
cv2.imwrite(temp_path, final)
|
|
|
|
try:
|
|
res = text_rec.predict(temp_path, batch_size=1)
|
|
texts = [r.get("rec_text", "") for r in res]
|
|
recognized = "".join(texts)
|
|
|
|
sim = similarity(recognized, TARGET_TEXT)
|
|
|
|
print(f"Testing {test_name}: '{recognized}' (Sim: {sim:.2f})")
|
|
|
|
if sim > max_sim:
|
|
max_sim = sim
|
|
found_best = (test_name, recognized, sim, temp_path)
|
|
|
|
if sim >= MIN_SIMILARITY:
|
|
print(f"SUCCESS! Found match with {test_name}")
|
|
# Keep going to find best, but we found a winner
|
|
else:
|
|
# Clean up failed images to save space
|
|
# os.remove(temp_path)
|
|
pass
|
|
except:
|
|
pass
|
|
|
|
print("\n--- Search Complete ---")
|
|
if found_best:
|
|
print(f"Best: {found_best[0]} | '{found_best[1]}' | Sim: {found_best[2]:.2f}")
|
|
print(f"Best Image: {found_best[3]}")
|
|
else:
|
|
print("No matches found above 0.0 similarity.")
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 3:
|
|
print("Usage: python brute_force_unwarp.py <image_path> <results_dir>")
|
|
sys.exit(1)
|
|
brute_force_search(sys.argv[1], sys.argv[2])
|