Update app.py

app.py

@@ -1,128 +1,144 @@
 import gradio as gr
 import cv2
 import numpy as np
+import matplotlib.pyplot as plt
 import json
 import math
-import matplotlib.pyplot as plt
 
-# === Helper Functions ===
+# === Helper: Rotated rectangle corners ===
 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)
+
+    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]
+    ])
+
     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) + np.array([cx, cy])
-    return rotated_corners.astype(np.float32)
 
+    rotated_corners = np.dot(local_corners, R.T)
+    corners = rotated_corners + np.array([cx, cy])
+    return corners.astype(np.float32)
+
+# === Preprocessing ===
 def preprocess_gray_clahe(img):
     gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-    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, detector_type, ratio_thresh=0.78):
-    if detector_type == "SIFT":
-        detector = cv2.SIFT_create(nfeatures=5000)
-        matcher = cv2.BFMatcher(cv2.NORM_L2)
-    elif detector_type == "BRISK":
-        detector = cv2.BRISK_create()
-        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
-    elif detector_type == "ORB":
-        detector = cv2.ORB_create(5000)
-        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
-    elif detector_type == "AKAZE":
-        detector = cv2.AKAZE_create()
-        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
-    elif detector_type == "KAZE":
-        detector = cv2.KAZE_create()
-        matcher = cv2.BFMatcher(cv2.NORM_L2)
+# === Feature detectors ===
+def get_detector(detector_name):
+    if detector_name == "SIFT":
+        return cv2.SIFT_create(nfeatures=5000)
+    elif detector_name == "ORB":
+        return cv2.ORB_create(5000)
+    elif detector_name == "BRISK":
+        return cv2.BRISK_create()
+    elif detector_name == "AKAZE":
+        return cv2.AKAZE_create()
+    elif detector_name == "KAZE":
+        return cv2.KAZE_create()
     else:
-        return None, None, []
+        return None
 
+def detect_and_match(img1_gray, img2_gray, detector_name, ratio_thresh=0.78):
+    detector = get_detector(detector_name)
     kp1, des1 = detector.detectAndCompute(img1_gray, None)
     kp2, des2 = detector.detectAndCompute(img2_gray, None)
-    if des1 is None or des2 is None:
-        return kp1, kp2, []
 
-    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
+    if detector_name in ["SIFT", "KAZE"]:
+        matcher = cv2.BFMatcher(cv2.NORM_L2)
+    else:
+        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
 
-def get_roi_points_from_json(json_file):
-    data = json.load(json_file)
-    area = data["printAreas"][0]
-    x = area["position"]["x"]
-    y = area["position"]["y"]
-    w = area["width"]
-    h = area["height"]
-    rot = area["rotation"]
-    return x, y, w, h, rot
+    matches = matcher.knnMatch(des1, des2, k=2)
+    good = []
+    for m, n in matches:
+        if m.distance < ratio_thresh * n.distance:
+            good.append(m)
+    return kp1, kp2, good
 
+# === Main processing ===
 def process_images(flat_img, persp_img, json_file):
-    # Preprocess
+    if flat_img is None or persp_img is None or json_file is None:
+        return [None]*6
+
+    # Load JSON
+    try:
+        data = json.load(open(json_file.name))
+    except Exception as e:
+        print("JSON read error:", e)
+        return [None]*6
+
+    roi = data["printAreas"][0]
+    roi_x = roi["position"]["x"]
+    roi_y = roi["position"]["y"]
+    roi_w = roi["width"]
+    roi_h = roi["height"]
+    roi_rot_deg = roi["rotation"]
+
+    # Preprocess images
     flat_gray = preprocess_gray_clahe(flat_img)
     persp_gray = preprocess_gray_clahe(persp_img)
-    x, y, w, h, rot = get_roi_points_from_json(json_file)
 
-    detectors = ["SIFT","BRISK","ORB","AKAZE","KAZE"]
-    gallery_images = []
+    detectors = ["SIFT", "ORB", "BRISK", "AKAZE", "KAZE"]
+    results = []
 
     for det in detectors:
-        kp1, kp2, matches = detect_and_match(flat_gray, persp_gray, det)
-        if len(matches) < 4:
-            # Skip if too few matches
+        kp1, kp2, good_matches = detect_and_match(flat_gray, persp_gray, det)
+        if len(good_matches) < 4:
+            print(f"Not enough matches for {det}")
+            results.append(None)
             continue
 
-        src_pts = np.float32([kp1[m.queryIdx].pt for m in matches]).reshape(-1,1,2)
-        dst_pts = np.float32([kp2[m.trainIdx].pt for m in matches]).reshape(-1,1,2)
-        H, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC,5.0)
-
-        # ROI in flat
-        roi_flat = get_rotated_rect_corners(x,y,w,h,rot)
-        flat_copy = flat_img.copy()
-        cv2.polylines(flat_copy, [roi_flat.astype(int)], True, (0,0,255),2)
-
-        # Project ROI to perspective
-        roi_persp = cv2.perspectiveTransform(roi_flat.reshape(-1,1,2), H).reshape(-1,2)
-        persp_copy = persp_img.copy()
-        cv2.polylines(persp_copy, [roi_persp.astype(int)], True, (0,255,0),2)
-        for px, py in roi_persp:
-            cv2.circle(persp_copy, (int(px),int(py)), 5, (255,0,0), -1)
-
-        # Side-by-side for this detector
-        fig, ax = plt.subplots(1,2,figsize=(12,6))
-        ax[0].imshow(flat_copy)
-        ax[0].set_title(f"Flat ROI - {det}")
-        ax[0].axis("off")
-        ax[1].imshow(persp_copy)
-        ax[1].set_title(f"Perspective ROI - {det}")
-        ax[1].axis("off")
-        plt.tight_layout()
-        filename = f"{det}_result.png"
-        plt.savefig(filename)
-        plt.close(fig)
-        gallery_images.append(filename)
+        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, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
 
+        # ROI corners
+        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)
+
+        # Draw ROI
+        flat_out = flat_img.copy()
+        persp_out = persp_img.copy()
+        cv2.polylines(flat_out, [roi_corners_flat.astype(int)], True, (255,0,0), 3)
+        cv2.polylines(persp_out, [roi_corners_persp.astype(int)], True, (0,255,0), 3)
+
+        results.append([flat_out, persp_out])
+
+    return results  # List of [flat_out, persp_out] for each detector
+
+# === Gradio Interface ===
+def wrap_gradio(flat_img, persp_img, json_file):
+    outputs = process_images(flat_img, persp_img, json_file)
+    # Flatten the outputs for Gallery display
+    gallery_images = []
+    for item in outputs:
+        if item is not None:
+            gallery_images.extend([item[0], item[1]])
     return gallery_images
 
 iface = gr.Interface(
-    fn=process_images,
+    fn=wrap_gradio,
     inputs=[
         gr.Image(type="numpy", label="Flat Image"),
         gr.Image(type="numpy", label="Perspective Image"),
         gr.File(type="filepath", label="JSON File")
-    ],  # <-- ye closing bracket should be ]
-    outputs=[  # <-- starts a new list
-        gr.Gallery(label="Results"),
-        gr.File(label="Download SIFT Result"),
-        gr.File(label="Download ORB Result"),
-        gr.File(label="Download BRISK Result"),
-        gr.File(label="Download AKAZE Result"),
-        gr.File(label="Download KAZE Result")
-    ],  # <-- should be ] not )
-    title="Homography & ROI Projection",
-    description="..."
+    ],
+    outputs=[
+        gr.Gallery(label="Results (Flat + Perspective per Detector)")
+    ],
+    title="Feature Detection with ROI Projection",
+    description="Shows SIFT, ORB, BRISK, AKAZE, KAZE feature-based ROI projections. Each detector outputs Flat and Perspective images."
 )
 
+iface.launch()