Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import json
 import gradio as gr
 
-# ---------------- Your Original Functions (Unchanged) ---------------- #
+# ---------------- 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)
@@ -58,7 +58,10 @@ def detect_and_match(img1_gray, img2_gray, method="SIFT", ratio_thresh=0.78):
         kp2, des2 = akaze.detectAndCompute(img2_gray, None)
         matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
     else:
-        return None, None, None
+        return None, None, []
+
+    if des1 is None or des2 is None:
+        return None, None, []
 
     raw_matches = matcher.knnMatch(des1, des2, k=2)
     good = []
@@ -67,16 +70,11 @@ def detect_and_match(img1_gray, img2_gray, method="SIFT", ratio_thresh=0.78):
             good.append(m)
     return kp1, kp2, good
 
-# ---------------- Processing Function for Gradio ---------------- #
-def homography_all_detectors(flat_img, persp_img, json_file):
-    if flat_img is None or persp_img is None:
-        return [None] * 6
-    
-    flat_bgr = cv2.cvtColor(flat_img, cv2.COLOR_RGB2BGR)
-    persp_bgr = cv2.cvtColor(persp_img, cv2.COLOR_RGB2BGR)
-
-    with open(json_file.name, 'r') as f:
-        mockup = json.load(f)
+# ---------------- Main Homography Function ----------------
+def homography_all_detectors(flat_file, persp_file, json_file):
+    flat_img = cv2.imread(flat_file.name)
+    persp_img = cv2.imread(persp_file.name)
+    mockup = json.load(open(json_file.name))
 
     roi_data = mockup["printAreas"][0]["position"]
     roi_x = roi_data["x"]
@@ -85,40 +83,18 @@ def homography_all_detectors(flat_img, persp_img, json_file):
     roi_h = mockup["printAreas"][0]["height"]
     roi_rot_deg = mockup["printAreas"][0]["rotation"]
 
-    flat_gray = preprocess_gray_clahe(flat_bgr)
-    persp_gray = preprocess_gray_clahe(persp_bgr)
+    flat_gray = preprocess_gray_clahe(flat_img)
+    persp_gray = preprocess_gray_clahe(persp_img)
 
-    detectors = ["SIFT", "ORB", "BRISK", "KAZE", "AKAZE"]
-    gallery_images = []
-    download_files = [None] * 5def homography_all_detectors(flat_img, persp_img, json_file):
-    if flat_img is None or persp_img is None:
-        return [None] * 6
-    
-    flat_bgr = cv2.cvtColor(flat_img, cv2.COLOR_RGB2BGR)
-    persp_bgr = cv2.cvtColor(persp_img, cv2.COLOR_RGB2BGR)
+    methods = ["SIFT", "ORB", "BRISK", "KAZE", "AKAZE"]
+    outputs = []
 
-    with open(json_file.name, 'r') as f:
-        mockup = json.load(f)
-
-    roi_data = mockup["printAreas"][0]["position"]
-    roi_x = roi_data["x"]
-    roi_y = roi_data["y"]
-    roi_w = mockup["printAreas"][0]["width"]
-    roi_h = mockup["printAreas"][0]["height"]
-    roi_rot_deg = mockup["printAreas"][0]["rotation"]
-
-    flat_gray = preprocess_gray_clahe(flat_bgr)
-    persp_gray = preprocess_gray_clahe(persp_bgr)
-
-    detectors = ["SIFT", "ORB", "BRISK", "KAZE", "AKAZE"]
-    gallery_images = []
-    download_files = [None] * 5
-
-    for i, method in enumerate(detectors):
+    for method in methods:
         kp1, kp2, good_matches = detect_and_match(flat_gray, persp_gray, method=method)
+
         if kp1 is None or kp2 is None or len(good_matches) < 4:
-            continue
-        
+            continue  # skip if no matches
+
         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)
@@ -129,68 +105,35 @@ def homography_all_detectors(flat_img, persp_img, json_file):
         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_debug = persp_bgr.copy()
+        persp_debug = persp_img.copy()
         cv2.polylines(persp_debug, [roi_corners_persp.astype(int)], True, (0,255,0), 2)
         for (px, py) in roi_corners_persp:
             cv2.circle(persp_debug, (int(px), int(py)), 5, (255,0,0), -1)
 
+        # Convert BGR -> RGB for display
         result_rgb = cv2.cvtColor(persp_debug, cv2.COLOR_BGR2RGB)
-        file_name = f"result_{method.lower()}.png"
-        cv2.imwrite(file_name, result_rgb[:, :, ::-1])  # save as BGR
 
-        gallery_images.append(result_rgb)   # ✅ only image
-        download_files[i] = file_name
+        outputs.append((result_rgb, f"{method} Result"))
 
-    return [gallery_images] + download_files
+    return outputs
 
+# ---------------- Gradio UI ----------------
+with gr.Blocks() as demo:
+    gr.Markdown("## Homography ROI Demo with Multiple Feature Detectors")
 
-    for i, method in enumerate(detectors):
-        kp1, kp2, good_matches = detect_and_match(flat_gray, persp_gray, method=method)
-        if kp1 is None or kp2 is None or len(good_matches) < 4:
-            continue
-        
-        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)
+    with gr.Row():
+        flat_input = gr.File(label="Upload Flat Image", file_types=[".jpg",".png",".jpeg"])
+        persp_input = gr.File(label="Upload Perspective Image", file_types=[".jpg",".png",".jpeg"])
+        json_input = gr.File(label="Upload mockup.json", file_types=[".json"])
 
-        if H is None:
-            continue
+    output_gallery = gr.Gallery(
+        label="Perspective ROI Results", 
+        columns=2, 
+        height=400, 
+        show_label=True
+    )
 
-        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)
+    run_btn = gr.Button("Run Homography")
+    run_btn.click(homography_all_detectors, inputs=[flat_input, persp_input, json_input], outputs=output_gallery)
 
-        persp_debug = persp_bgr.copy()
-        cv2.polylines(persp_debug, [roi_corners_persp.astype(int)], True, (0,255,0), 2)
-        for (px, py) in roi_corners_persp:
-            cv2.circle(persp_debug, (int(px), int(py)), 5, (255,0,0), -1)
-
-        result_rgb = cv2.cvtColor(persp_debug, cv2.COLOR_BGR2RGB)
-        file_name = f"result_{method.lower()}.png"
-        cv2.imwrite(file_name, result_rgb[:, :, ::-1])  # save as BGR
-
-        gallery_images.append(result_rgb)   # ✅ only image
-        download_files[i] = file_name
-
-    return [gallery_images] + download_files
-
-# ---------------- Gradio Interface ---------------- #
-iface = gr.Interface(
-    fn=homography_all_detectors,
-    inputs=[
-        gr.Image(type="numpy", label="Image 1 (Flat)"),
-        gr.Image(type="numpy", label="Image 2 (Perspective)"),
-        gr.File(type="filepath", label="JSON File")
-    ],
-    outputs=[
-        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 KAZE Result"),
-        gr.File(label="Download AKAZE Result")
-    ],
-    title="Homography ROI Projection with Multiple Feature Detectors",
-    description="Upload a flat image, a perspective image, and the JSON file. The system will compute homography with SIFT, ORB, BRISK, KAZE, and AKAZE, project the bounding box, and allow result download."
-)
-
-iface.launch()
+demo.launch()