Update app.py

app.py

@@ -5,225 +5,139 @@ import gradio as gr
 import os
 import xml.etree.ElementTree as ET
 
-# ---------------- Helper functions ----------------
+# ---------------- Helper Functions ----------------
 def get_rotated_rect_corners(x, y, w, h, rotation_deg):
     rot_rad = np.deg2rad(rotation_deg)
-    cos_r = np.cos(rot_rad)
-    sin_r = np.sin(rot_rad)
-    R = np.array([[cos_r, -sin_r],
-                  [sin_r,  cos_r]])
-    cx = x + w/2
-    cy = y + h/2
-    local_corners = np.array([
-        [-w/2, -h/2],
-        [ w/2, -h/2],
-        [ w/2,  h/2],
-        [-w/2,  h/2]
-    ])
+    cos_r, sin_r = np.cos(rot_rad), np.sin(rot_rad)
+    R = np.array([[cos_r, -sin_r], [sin_r, cos_r]])
+    cx, cy = x + w/2, y + h/2
+    local_corners = np.array([[-w/2,-h/2],[w/2,-h/2],[w/2,h/2],[-w/2,h/2]])
     rotated_corners = np.dot(local_corners, R.T)
-    corners = rotated_corners + np.array([cx, cy])
-    return corners.astype(np.float32)
+    return (rotated_corners + np.array([cx,cy])).astype(np.float32)
 
 def preprocess_gray_clahe(img):
     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
+    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
     return clahe.apply(gray)
 
 def detect_and_match(img1_gray, img2_gray, method="SIFT", ratio_thresh=0.78):
-    if method == "SIFT":
-        detector = cv2.SIFT_create(nfeatures=5000)
-        norm = cv2.NORM_L2
-    elif method == "ORB":
-        detector = cv2.ORB_create(5000)
-        norm = cv2.NORM_HAMMING
-    elif method == "BRISK":
-        detector = cv2.BRISK_create()
-        norm = cv2.NORM_HAMMING
-    elif method == "KAZE":
-        detector = cv2.KAZE_create()
-        norm = cv2.NORM_L2
-    elif method == "AKAZE":
-        detector = cv2.AKAZE_create()
-        norm = cv2.NORM_HAMMING
-    else:
-        return None, None, [], None
-
-    kp1, des1 = detector.detectAndCompute(img1_gray, None)
-    kp2, des2 = detector.detectAndCompute(img2_gray, None)
-    if des1 is None or des2 is None:
-        return None, None, [], None
-
-    matcher = cv2.BFMatcher(norm)
-    raw_matches = matcher.knnMatch(des1, des2, k=2)
-    good = [m for m,n in raw_matches if m.distance < ratio_thresh * n.distance]
-
-    matches_img = None
-    if len(good) >= 4:
-        matches_img = cv2.drawMatches(
-            cv2.cvtColor(img1_gray, cv2.COLOR_GRAY2BGR),
-            kp1,
-            cv2.cvtColor(img2_gray, cv2.COLOR_GRAY2BGR),
-            kp2,
-            good, None,
-            flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS
-        )
-    return kp1, kp2, good, matches_img
-
-def add_title(img_bgr, title):
-    h, w = img_bgr.shape[:2]
-    bar = np.full((40, w, 3), 255, dtype=np.uint8)
-    cv2.putText(bar, title, (10, 28), cv2.FONT_HERSHEY_SIMPLEX,
-                0.8, (0,0,0), 2, cv2.LINE_AA)
-    return np.vstack([bar, img_bgr])
-
-def resize_to_height(img, target_height):
-    """Resize image to target height while maintaining aspect ratio"""
-    h, w = img.shape[:2]
-    ratio = target_height / h
-    new_width = int(w * ratio)
-    return cv2.resize(img, (new_width, target_height))
-
-def parse_xml_roi_points(xml_path):
-    """Parse your XML structure, return list of polygons (Nx2 points)."""
-    if xml_path is None:
-        return None
-        
-    polys = []
-    try:
-        tree = ET.parse(xml_path)
-        root = tree.getroot()
-        # Transform ROI points (FourPoint)
-        for tr in root.findall(".//transform"):
-            pts = []
-            for p in tr.findall("point"):
-                x = float(p.get("x")); y = float(p.get("y"))
-                pts.append([x, y])
-            if len(pts) >= 3:
-                polys.append(np.array(pts, dtype=np.float32))
-        # Overlay polygons
-        for ov in root.findall(".//overlay"):
-            pts = []
-            for p in ov.findall("point"):
-                x = float(p.get("x")); y = float(p.get("y"))
-                pts.append([x, y])
-            if len(pts) >= 3:
-                polys.append(np.array(pts, dtype=np.float32))
-    except Exception as e:
-        print("XML parse error:", e)
-    return polys if polys else None
+    if method=="SIFT": detector=cv2.SIFT_create(nfeatures=5000); matcher=cv2.BFMatcher(cv2.NORM_L2)
+    elif method=="ORB": detector=cv2.ORB_create(5000); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
+    elif method=="BRISK": detector=cv2.BRISK_create(); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
+    elif method=="KAZE": detector=cv2.KAZE_create(); matcher=cv2.BFMatcher(cv2.NORM_L2)
+    elif method=="AKAZE": detector=cv2.AKAZE_create(); matcher=cv2.BFMatcher(cv2.NORM_HAMMING)
+    else: return None,None,[]
+    
+    kp1, des1 = detector.detectAndCompute(img1_gray,None)
+    kp2, des2 = detector.detectAndCompute(img2_gray,None)
+    if des1 is None or des2 is None: return None,None,[]
+    
+    raw_matches = matcher.knnMatch(des1,des2,k=2)
+    good = [m for m,n in raw_matches if m.distance < ratio_thresh*n.distance]
+    return kp1, kp2, good
+
+def parse_xml_points(xml_file):
+    tree = ET.parse(xml_file)
+    root = tree.getroot()
+    points=[]
+    for pt_type in ["TopLeft","TopRight","BottomLeft","BottomRight"]:
+        elem=root.find(f".//point[@type='{pt_type}']")
+        points.append([float(elem.get("x")), float(elem.get("y"))])
+    return np.array(points,dtype=np.float32).reshape(-1,2)
 
 # ---------------- Main Function ----------------
 def homography_all_detectors(flat_file, persp_file, json_file, xml_file):
     flat_img = cv2.imread(flat_file)
     persp_img = cv2.imread(persp_file)
-    
-    # Use the file paths directly (no .name needed)
-    mockup = json.load(open(json_file))
-    roi = mockup["printAreas"][0]
-    roi_x, roi_y = roi["position"]["x"], roi["position"]["y"]
-    roi_w, roi_h = roi["width"], roi["height"]
-    roi_rot_deg = roi["rotation"]
-
-    xml_polys = parse_xml_roi_points(xml_file) if xml_file else None
+    mockup = json.load(open(json_file.name))
+    roi_data = mockup["printAreas"][0]["position"]
+    roi_x, roi_y = roi_data["x"], roi_data["y"]
+    roi_w, roi_h = mockup["printAreas"][0]["width"], mockup["printAreas"][0]["height"]
+    roi_rot_deg = mockup["printAreas"][0]["rotation"]
 
     flat_gray = preprocess_gray_clahe(flat_img)
     persp_gray = preprocess_gray_clahe(persp_img)
+    xml_points = parse_xml_points(xml_file.name)
 
-    results = []
+    methods = ["SIFT","ORB","BRISK","KAZE","AKAZE"]
+    gallery_images = []
     download_files = []
 
-    for method in ["SIFT","ORB","BRISK","KAZE","AKAZE"]:
-        try:
-            kp1, kp2, good, matches_img = detect_and_match(flat_gray, persp_gray, method)
-            if kp1 is None or len(good) < 4:
-                continue
-
-            src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1,1,2)
-            dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1,1,2)
-            H, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
-
-            # top-left = flat
-            top_left = add_title(flat_img.copy(), "Flat Image")
-
-            # top-right = feature matches
-            if matches_img is None:
-                matches_img = flat_img.copy()
-            top_right = add_title(matches_img, "Feature Matches (Flat→Perspective)")
-
-            # bottom-left = homography ROI
-            bottom_left = persp_img.copy()
-            if H is not None:
-                roi_corners = get_rotated_rect_corners(roi_x, roi_y, roi_w, roi_h, roi_rot_deg)
-                roi_persp = cv2.perspectiveTransform(roi_corners.reshape(-1,1,2), H).reshape(-1,2)
-                cv2.polylines(bottom_left, [roi_persp.astype(int)], True, (0,255,0), 3)
-            bottom_left = add_title(bottom_left, "ROI via Homography (from JSON)")
-
-            # bottom-right = XML ROI
-            bottom_right = persp_img.copy()
-            if xml_polys:
-                for poly in xml_polys:
-                    cv2.polylines(bottom_right, [poly.astype(int)], True, (0,0,255), 3)
-            else:
-                cv2.putText(bottom_right, "No XML ROI", (50,100),
-                            cv2.FONT_HERSHEY_SIMPLEX, 2, (0,0,255), 3)
-            bottom_right = add_title(bottom_right, "Ground Truth ROI (XML)")
-
-            # Resize all images to the same height before stacking
-            target_height = max(top_left.shape[0], top_right.shape[0], 
-                               bottom_left.shape[0], bottom_right.shape[0])
-            
-            top_left_resized = resize_to_height(top_left, target_height)
-            top_right_resized = resize_to_height(top_right, target_height)
-            bottom_left_resized = resize_to_height(bottom_left, target_height)
-            bottom_right_resized = resize_to_height(bottom_right, target_height)
-
-            # stack horizontally
-            top = np.hstack([top_left_resized, top_right_resized])
-            bottom = np.hstack([bottom_left_resized, bottom_right_resized])
-            
-            # Make sure top and bottom have the same width
-            min_width = min(top.shape[1], bottom.shape[1])
-            top = cv2.resize(top, (min_width, top.shape[0]))
-            bottom = cv2.resize(bottom, (min_width, bottom.shape[0]))
-            
-            composite = np.vstack([top, bottom])
-
-            base_name = os.path.splitext(os.path.basename(persp_file))[0]
-            file_name = f"{base_name}_{method.lower()}_grid.png"
-            cv2.imwrite(file_name, composite)
-            results.append((cv2.cvtColor(composite, cv2.COLOR_BGR2RGB), f"{method} Grid"))
-            download_files.append(file_name)
-        except Exception as e:
-            print(f"Error in {method}: {str(e)}")
-            continue
-
-    while len(download_files) < 5:
-        download_files.append(None)
-    
-    # Return the results in the correct format
-    gallery_output = results
-    file_outputs = download_files[:5]
-    
-    return [gallery_output] + file_outputs
+    for method in methods:
+        kp1,kp2,good_matches = detect_and_match(flat_gray,persp_gray,method)
+        if kp1 is None or kp2 is None or len(good_matches)<4: continue
+
+        # Feature matching
+        match_img = cv2.drawMatches(flat_img,kp1,persp_img,kp2,good_matches,None,flags=2)
+
+        # Homography & ROI
+        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1,1,2)
+        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1,1,2)
+        H,_ = cv2.findHomography(src_pts,dst_pts,cv2.RANSAC,5.0)
+        if H is None: continue
+
+        roi_corners_flat = get_rotated_rect_corners(roi_x,roi_y,roi_w,roi_h,roi_rot_deg)
+        roi_corners_persp = cv2.perspectiveTransform(roi_corners_flat.reshape(-1,1,2),H).reshape(-1,2)
+        persp_roi = persp_img.copy()
+        cv2.polylines(persp_roi,[roi_corners_persp.astype(int)],True,(0,255,0),2)
+        for px,py in roi_corners_persp: cv2.circle(persp_roi,(int(px),int(py)),5,(255,0,0),-1)
+
+        xml_gt_img = persp_img.copy()
+        xml_mapped = cv2.perspectiveTransform(xml_points.reshape(-1,1,2),H).reshape(-1,2)
+        for px,py in xml_mapped: cv2.circle(xml_gt_img,(int(px),int(py)),5,(0,0,255),-1)
+
+        # --------- For display only: resize to match grid ---------
+        def resize_to_height(img, target_h):
+            h, w = img.shape[:2]
+            scale = target_h / h
+            new_w = int(w * scale)
+            return cv2.resize(img, (new_w, target_h))
+
+        target_h = 300  # temporary display height for gallery
+        flat_disp = resize_to_height(flat_img, target_h)
+        match_disp = resize_to_height(match_img, target_h)
+        roi_disp = resize_to_height(persp_roi, target_h)
+        xml_disp = resize_to_height(xml_gt_img, target_h)
+
+        # Merge 2x2 grid for gallery display
+        top = np.hstack([flat_disp, match_disp])
+        bottom = np.hstack([roi_disp, xml_disp])
+        combined_grid = np.vstack([top, bottom])
+        gallery_images.append((combined_grid,f"{method} Detector"))
+
+        # Save original resolution for download
+        base_name = os.path.splitext(os.path.basename(persp_file))[0]
+        file_name = f"{base_name}_{method.lower()}.png"
+        # Merge original images for download
+        top_orig = np.hstack([flat_img, match_img])
+        bottom_orig = np.hstack([persp_roi, xml_gt_img])
+        combined_orig = np.vstack([top_orig, bottom_orig])
+        cv2.imwrite(file_name, combined_orig)
+        download_files.append(file_name)
+
+    # Ensure 5 outputs
+    while len(download_files)<5: download_files.append(None)
+    return gallery_images, download_files[0], download_files[1], download_files[2], download_files[3], download_files[4]
 
 # ---------------- Gradio UI ----------------
 iface = gr.Interface(
     fn=homography_all_detectors,
     inputs=[
-        gr.Image(label="Upload Flat Image", type="filepath"),
-        gr.Image(label="Upload Perspective Image", type="filepath"),
-        gr.File(label="Upload mockup.json", file_types=[".json"]),
-        gr.File(label="Upload Groundtruth.xml", file_types=[".xml"])
+        gr.Image(label="Upload Flat Image",type="filepath"),
+        gr.Image(label="Upload Perspective Image",type="filepath"),
+        gr.File(label="Upload mockup.json",file_types=[".json"]),
+        gr.File(label="Upload XML file",file_types=[".xml"])
     ],
     outputs=[
-        gr.Gallery(label="Composite Grids (per Detector)", show_label=True),
-        gr.File(label="Download SIFT Grid"),
-        gr.File(label="Download ORB Grid"),
-        gr.File(label="Download BRISK Grid"),
-        gr.File(label="Download KAZE Grid"),
-        gr.File(label="Download AKAZE Grid")
+        gr.Gallery(label="Results per Detector",show_label=True),
+        gr.File(label="Download SIFT Result"),
+        gr.File(label="Download ORB Result"),
+        gr.File(label="Download BRISK Result"),
+        gr.File(label="Download KAZE Result"),
+        gr.File(label="Download AKAZE Result")
     ],
-    title="Homography ROI vs XML ROI",
-    description="Each detector produces one 2×2 grid: Flat, Matches, Homography ROI, Ground Truth ROI."
+    title="Homography ROI Projection with Feature Matching & XML GT",
+    description="Shows 4 views per detector. Grid resized for display, original resolution saved for download."
 )
-iface.launch()
+
+iface.launch()