Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import numpy as np
|
3 |
+
import json
|
4 |
+
import gradio as gr
|
5 |
+
import os
|
6 |
+
import xml.etree.ElementTree as ET
|
7 |
+
|
8 |
+
# ---------------- Helper functions ----------------
|
9 |
+
def get_rotated_rect_corners(x, y, w, h, rotation_deg):
|
10 |
+
rot_rad = np.deg2rad(rotation_deg)
|
11 |
+
cos_r, sin_r = np.cos(rot_rad), np.sin(rot_rad)
|
12 |
+
R = np.array([[cos_r, -sin_r], [sin_r, cos_r]])
|
13 |
+
cx, cy = x + w/2, y + h/2
|
14 |
+
local_corners = np.array([[-w/2,-h/2],[w/2,-h/2],[w/2,h/2],[-w/2,h/2]])
|
15 |
+
rotated_corners = np.dot(local_corners, R.T)
|
16 |
+
return (rotated_corners + np.array([cx,cy])).astype(np.float32)
|
17 |
+
|
18 |
+
def preprocess_gray_clahe(img):
|
19 |
+
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
20 |
+
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
|
21 |
+
return clahe.apply(gray)
|
22 |
+
|
23 |
+
def detect_and_match(img1_gray, img2_gray, method="SIFT", ratio_thresh=0.78):
|
24 |
+
if method=="SIFT": detector=cv2.SIFT_create(nfeatures=5000); matcher=cv2.BFMatcher(cv2.NORM_L2)
|
25 |
+
elif method=="ORB": detector=cv2.ORB_create(5000); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
|
26 |
+
elif method=="BRISK": detector=cv2.BRISK_create(); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
|
27 |
+
elif method=="KAZE": detector=cv2.KAZE_create(); matcher=cv2.BFMatcher(cv2.NORM_L2)
|
28 |
+
elif method=="AKAZE": detector=cv2.AKAZE_create(); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
|
29 |
+
else: return None,None,[]
|
30 |
+
|
31 |
+
kp1, des1 = detector.detectAndCompute(img1_gray,None)
|
32 |
+
kp2, des2 = detector.detectAndCompute(img2_gray,None)
|
33 |
+
if des1 is None or des2 is None: return None,None,[]
|
34 |
+
|
35 |
+
raw_matches = matcher.knnMatch(des1,des2,k=2)
|
36 |
+
good = [m for m,n in raw_matches if m.distance < ratio_thresh*n.distance]
|
37 |
+
return kp1, kp2, good
|
38 |
+
|
39 |
+
def parse_xml_points(xml_file):
|
40 |
+
tree = ET.parse(xml_file)
|
41 |
+
root = tree.getroot()
|
42 |
+
points=[]
|
43 |
+
for pt_type in ["TopLeft","TopRight","BottomLeft","BottomRight"]:
|
44 |
+
elem=root.find(f".//point[@type='{pt_type}']")
|
45 |
+
points.append([float(elem.get("x")), float(elem.get("y"))])
|
46 |
+
return np.array(points,dtype=np.float32).reshape(-1,2)
|
47 |
+
|
48 |
+
# ---------------- Padding Helper ----------------
|
49 |
+
def pad_to_size(img, target_h, target_w):
|
50 |
+
h, w = img.shape[:2]
|
51 |
+
top_pad = 0
|
52 |
+
left_pad = 0
|
53 |
+
bottom_pad = target_h - h
|
54 |
+
right_pad = target_w - w
|
55 |
+
canvas = np.ones((target_h, target_w,3), dtype=np.uint8)*255
|
56 |
+
canvas[top_pad:top_pad+h, left_pad:left_pad+w] = img
|
57 |
+
return canvas
|
58 |
+
|
59 |
+
# ---------------- Resize feature-match to original reference size ----------------
|
60 |
+
def match_img_to_reference(match_img, ref_h, ref_w):
|
61 |
+
h, w = match_img.shape[:2]
|
62 |
+
scale = min(ref_w/w, ref_h/h)
|
63 |
+
new_w, new_h = int(w*scale), int(h*scale)
|
64 |
+
resized = cv2.resize(match_img, (new_w,new_h))
|
65 |
+
padded = pad_to_size(resized, ref_h, ref_w)
|
66 |
+
return padded
|
67 |
+
|
68 |
+
# ---------------- Main Function ----------------
|
69 |
+
def homography_all_detectors(flat_file, persp_file, json_file, xml_file):
|
70 |
+
flat_img = cv2.imread(flat_file)
|
71 |
+
persp_img = cv2.imread(persp_file)
|
72 |
+
mockup = json.load(open(json_file.name))
|
73 |
+
roi_data = mockup["printAreas"][0]["position"]
|
74 |
+
roi_x, roi_y = roi_data["x"], roi_data["y"]
|
75 |
+
roi_w, roi_h = mockup["printAreas"][0]["width"], mockup["printAreas"][0]["height"]
|
76 |
+
roi_rot_deg = mockup["printAreas"][0]["rotation"]
|
77 |
+
|
78 |
+
flat_gray = preprocess_gray_clahe(flat_img)
|
79 |
+
persp_gray = preprocess_gray_clahe(persp_img)
|
80 |
+
xml_points = parse_xml_points(xml_file.name)
|
81 |
+
|
82 |
+
methods = ["SIFT","ORB","BRISK","KAZE","AKAZE"]
|
83 |
+
gallery_paths = []
|
84 |
+
download_files = []
|
85 |
+
|
86 |
+
for method in methods:
|
87 |
+
kp1,kp2,good_matches = detect_and_match(flat_gray,persp_gray,method)
|
88 |
+
if kp1 is None or kp2 is None or len(good_matches)<4: continue
|
89 |
+
|
90 |
+
match_img = cv2.drawMatches(flat_img,kp1,persp_img,kp2,good_matches,None,flags=2)
|
91 |
+
|
92 |
+
src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1,1,2)
|
93 |
+
dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1,1,2)
|
94 |
+
H,_ = cv2.findHomography(src_pts,dst_pts,cv2.RANSAC,5.0)
|
95 |
+
if H is None: continue
|
96 |
+
|
97 |
+
roi_corners_flat = get_rotated_rect_corners(roi_x,roi_y,roi_w,roi_h,roi_rot_deg)
|
98 |
+
roi_corners_persp = cv2.perspectiveTransform(roi_corners_flat.reshape(-1,1,2),H).reshape(-1,2)
|
99 |
+
persp_roi = persp_img.copy()
|
100 |
+
cv2.polylines(persp_roi,[roi_corners_persp.astype(int)],True,(0,255,0),2)
|
101 |
+
for px,py in roi_corners_persp: cv2.circle(persp_roi,(int(px),int(py)),5,(255,0,0),-1)
|
102 |
+
|
103 |
+
xml_gt_img = persp_img.copy()
|
104 |
+
xml_mapped = cv2.perspectiveTransform(xml_points.reshape(-1,1,2),H).reshape(-1,2)
|
105 |
+
for px,py in xml_mapped: cv2.circle(xml_gt_img,(int(px),int(py)),5,(0,0,255),-1)
|
106 |
+
|
107 |
+
# Convert to RGB
|
108 |
+
flat_rgb = cv2.cvtColor(flat_img,cv2.COLOR_BGR2RGB)
|
109 |
+
persp_rgb = cv2.cvtColor(persp_img,cv2.COLOR_BGR2RGB)
|
110 |
+
roi_rgb = cv2.cvtColor(persp_roi,cv2.COLOR_BGR2RGB)
|
111 |
+
xml_rgb = cv2.cvtColor(xml_gt_img,cv2.COLOR_BGR2RGB)
|
112 |
+
|
113 |
+
# Resize feature-match image to match original flat/perspective
|
114 |
+
match_rgb = match_img_to_reference(cv2.cvtColor(match_img, cv2.COLOR_BGR2RGB), flat_rgb.shape[0], flat_rgb.shape[1])
|
115 |
+
|
116 |
+
# Determine max height and width for grid (all images now same)
|
117 |
+
max_h = max(flat_rgb.shape[0], match_rgb.shape[0], roi_rgb.shape[0], xml_rgb.shape[0])
|
118 |
+
max_w = max(flat_rgb.shape[1], match_rgb.shape[1], roi_rgb.shape[1], xml_rgb.shape[1])
|
119 |
+
|
120 |
+
flat_pad = pad_to_size(flat_rgb, max_h, max_w)
|
121 |
+
roi_pad = pad_to_size(roi_rgb, max_h, max_w)
|
122 |
+
xml_pad = pad_to_size(xml_rgb, max_h, max_w)
|
123 |
+
|
124 |
+
# Merge 2x2 grid
|
125 |
+
top = np.hstack([flat_pad, match_rgb])
|
126 |
+
bottom = np.hstack([roi_pad, xml_pad])
|
127 |
+
combined_grid = np.vstack([top, bottom])
|
128 |
+
|
129 |
+
base_name = os.path.splitext(os.path.basename(persp_file))[0]
|
130 |
+
file_name = f"{base_name}_{method.lower()}.png"
|
131 |
+
cv2.imwrite(file_name, cv2.cvtColor(combined_grid,cv2.COLOR_RGB2BGR))
|
132 |
+
gallery_paths.append(file_name)
|
133 |
+
download_files.append(file_name)
|
134 |
+
|
135 |
+
while len(download_files)<5: download_files.append(None)
|
136 |
+
return gallery_paths, download_files[0], download_files[1], download_files[2], download_files[3], download_files[4]
|
137 |
+
|
138 |
+
iface = gr.Interface(
|
139 |
+
fn=homography_all_detectors,
|
140 |
+
inputs=[
|
141 |
+
gr.Image(label="Upload Flat Image",type="filepath"),
|
142 |
+
gr.Image(label="Upload Perspective Image",type="filepath"),
|
143 |
+
gr.File(label="Upload mockup.json",file_types=[".json"]),
|
144 |
+
gr.File(label="Upload XML file",file_types=[".xml"])
|
145 |
+
],
|
146 |
+
outputs=[
|
147 |
+
gr.Gallery(label="Results per Detector",show_label=True),
|
148 |
+
gr.File(label="Download SIFT Result"),
|
149 |
+
gr.File(label="Download ORB Result"),
|
150 |
+
gr.File(label="Download BRISK Result"),
|
151 |
+
gr.File(label="Download KAZE Result"),
|
152 |
+
gr.File(label="Download AKAZE Result")
|
153 |
+
],
|
154 |
+
title="Homography ROI Projection with Feature Matching & XML GT",
|
155 |
+
description="Flat + Perspective images with mockup.json & XML. Feature-match aligned with original images using white padding."
|
156 |
+
)
|
157 |
+
|
158 |
+
iface.launch()
|