Update app.py

app.py

@@ -4,7 +4,6 @@ import json
 import gradio as gr
 import os
 import xml.etree.ElementTree as ET
-import traceback
 
 # ---------------- Helper functions ----------------
 def get_rotated_rect_corners(x, y, w, h, rotation_deg):
@@ -32,15 +31,20 @@ def preprocess_gray_clahe(img):
 
 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
+        detector = cv2.SIFT_create(nfeatures=5000)
+        norm = cv2.NORM_L2
     elif method == "ORB":
-        detector = cv2.ORB_create(5000); norm = cv2.NORM_HAMMING
+        detector = cv2.ORB_create(5000)
+        norm = cv2.NORM_HAMMING
     elif method == "BRISK":
-        detector = cv2.BRISK_create(); norm = cv2.NORM_HAMMING
+        detector = cv2.BRISK_create()
+        norm = cv2.NORM_HAMMING
     elif method == "KAZE":
-        detector = cv2.KAZE_create(); norm = cv2.NORM_L2
+        detector = cv2.KAZE_create()
+        norm = cv2.NORM_L2
     elif method == "AKAZE":
-        detector = cv2.AKAZE_create(); norm = cv2.NORM_HAMMING
+        detector = cv2.AKAZE_create()
+        norm = cv2.NORM_HAMMING
     else:
         return None, None, [], None
 
@@ -56,8 +60,10 @@ def detect_and_match(img1_gray, img2_gray, method="SIFT", ratio_thresh=0.78):
     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,
+            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
         )
@@ -70,127 +76,135 @@ def add_title(img_bgr, title):
                 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 = []
-    tree = ET.parse(xml_path)
-    root = tree.getroot()
-    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))
-    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))
+    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
 
-def make_error_image(msg, width=800, height=600):
-    img = np.ones((height, width, 3), dtype=np.uint8)*255
-    y0 = 100
-    for i, line in enumerate(msg.split("\n")):
-        y = y0 + i*40
-        cv2.putText(img, line, (50, y), cv2.FONT_HERSHEY_SIMPLEX,
-                    0.8, (0,0,255), 2, cv2.LINE_AA)
-    return img
-
 # ---------------- 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
+
+    flat_gray = preprocess_gray_clahe(flat_img)
+    persp_gray = preprocess_gray_clahe(persp_img)
+
     results = []
     download_files = []
-    methods = ["SIFT","ORB","BRISK","KAZE","AKAZE"]
 
-    try:
-        flat_img = cv2.imread(flat_file)
-        persp_img = cv2.imread(persp_file)
-        if flat_img is None or persp_img is None:
-            raise ValueError("Could not read input images")
-
-        mockup = json.load(open(json_file.name))
-        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.name) if xml_file else None
-
-        flat_gray = preprocess_gray_clahe(flat_img)
-        persp_gray = preprocess_gray_clahe(persp_img)
-
-        for method in methods:
-            try:
-                kp1, kp2, good, matches_img = detect_and_match(flat_gray, persp_gray, method)
-                if kp1 is None or len(good) < 4:
-                    raise ValueError(f"{method}: Not enough good matches")
-
-                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)
-                if H is None:
-                    raise ValueError(f"{method}: Homography could not be estimated")
-
-                # top-left = flat
-                top_left = add_title(flat_img.copy(), "Flat Image")
-                # top-right = matches
-                top_right = add_title(matches_img if matches_img is not None else flat_img.copy(),
-                                      "Feature Matches")
-                # bottom-left = homography ROI
-                bottom_left = persp_img.copy()
+    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 (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)")
-
-                top = np.hstack([top_left, top_right])
-                bottom = np.hstack([bottom_left, bottom_right])
-                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:
-                err_msg = f"Error in {method}: {str(e)}"
-                print(err_msg)
-                tb = traceback.format_exc()
-                err_img = make_error_image(err_msg + "\n" + tb)
-                results.append((cv2.cvtColor(err_img, cv2.COLOR_BGR2RGB), f"{method} ERROR"))
-                # make txt file also
-                err_file = f"{method.lower()}_error.txt"
-                with open(err_file, "w") as f:
-                    f.write(err_msg + "\n\n" + tb)
-                download_files.append(err_file)
-
-    except Exception as e:
-        err_msg = f"Global Error: {str(e)}"
-        tb = traceback.format_exc()
-        err_img = make_error_image(err_msg + "\n" + tb)
-        results.append((cv2.cvtColor(err_img, cv2.COLOR_BGR2RGB), f"GLOBAL ERROR"))
-        err_file = "global_error.txt"
-        with open(err_file, "w") as f:
-            f.write(err_msg + "\n\n" + tb)
-        download_files = [err_file]*5
+            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 [results] + download_files[:5]
+    
+    # Return the results in the correct format
+    gallery_output = results
+    file_outputs = download_files[:5]
+    
+    return [gallery_output] + file_outputs
 
 # ---------------- Gradio UI ----------------
 iface = gr.Interface(
@@ -210,7 +224,6 @@ iface = gr.Interface(
         gr.File(label="Download AKAZE Grid")
     ],
     title="Homography ROI vs XML ROI",
-    description="Each detector produces one 2×2 grid: Flat, Matches, Homography ROI, Ground Truth ROI. If any error occurs, it will be shown in the output image + downloadable txt log."
+    description="Each detector produces one 2×2 grid: Flat, Matches, Homography ROI, Ground Truth ROI."
 )
-
-iface.launch()
+iface.launch()