Update app.py

app.py

@@ -1,135 +1,128 @@
-import gradio as gr
 import cv2
 import numpy as np
 import json
-import math
+import gradio as gr
+import matplotlib.pyplot as plt
 
 # === 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)
+
+    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]
     ])
-    R = np.array([[cos_r, -sin_r],
-                  [sin_r,  cos_r]])
+
     rotated_corners = np.dot(local_corners, R.T)
     corners = rotated_corners + np.array([cx, cy])
     return corners.astype(np.float32)
 
-# === Preprocessing ===
+# === Preprocessing: Grayscale + CLAHE ===
 def preprocess_gray_clahe(img):
-    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
     return clahe.apply(gray)
 
-# === 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
-
-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 detector_name in ["SIFT", "KAZE"]:
+# === Detect and match features ===
+def detect_and_match(img1_gray, img2_gray, method="SIFT", ratio_thresh=0.78):
+    if method == "SIFT":
+        sift = cv2.SIFT_create(nfeatures=5000)
+        kp1, des1 = sift.detectAndCompute(img1_gray, None)
+        kp2, des2 = sift.detectAndCompute(img2_gray, None)
         matcher = cv2.BFMatcher(cv2.NORM_L2)
-    else:
+    elif method == "ORB":
+        orb = cv2.ORB_create(5000)
+        kp1, des1 = orb.detectAndCompute(img1_gray, None)
+        kp2, des2 = orb.detectAndCompute(img2_gray, None)
+        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
+    elif method == "BRISK":
+        brisk = cv2.BRISK_create()
+        kp1, des1 = brisk.detectAndCompute(img1_gray, None)
+        kp2, des2 = brisk.detectAndCompute(img2_gray, None)
         matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
-    matches = matcher.knnMatch(des1, des2, k=2)
-    good = [m for m,n in matches if m.distance < ratio_thresh*n.distance]
+    elif method == "KAZE":
+        kaze = cv2.KAZE_create()
+        kp1, des1 = kaze.detectAndCompute(img1_gray, None)
+        kp2, des2 = kaze.detectAndCompute(img2_gray, None)
+        matcher = cv2.BFMatcher(cv2.NORM_L2)
+    elif method == "AKAZE":
+        akaze = cv2.AKAZE_create()
+        kp1, des1 = akaze.detectAndCompute(img1_gray, None)
+        kp2, des2 = akaze.detectAndCompute(img2_gray, None)
+        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
+    else:
+        return None, None, None
+
+    raw_matches = matcher.knnMatch(des1, des2, k=2)
+    good = []
+    for m, n in raw_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):
-    if flat_img is None or persp_img is None or json_file is None:
-        return []
-
-    # Load JSON
-    try:
-        with open(json_file.name, 'r') as f:
-            data = json.load(f)
-    except Exception as e:
-        print("JSON read error:", e)
-        return []
-
-    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
+# === Main function for Gradio ===
+def homography_demo(flat_file, persp_file, json_file):
+    flat_img = cv2.imdecode(np.frombuffer(flat_file.read(), np.uint8), cv2.IMREAD_COLOR)
+    persp_img = cv2.imdecode(np.frombuffer(persp_file.read(), np.uint8), cv2.IMREAD_COLOR)
+    mockup = json.load(json_file)
+
+    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_img)
     persp_gray = preprocess_gray_clahe(persp_img)
 
-    detectors = ["SIFT", "ORB", "BRISK", "AKAZE", "KAZE"]
-    gallery_images = []
+    methods = ["SIFT", "ORB", "BRISK", "KAZE", "AKAZE"]
+    outputs = []
 
-    for det in detectors:
-        kp1, kp2, good_matches = detect_and_match(flat_gray, persp_gray, det)
-        if len(good_matches) < 4:
-            print(f"Not enough matches for {det}")
-            continue
+    for method in methods:
+        kp1, kp2, good_matches = detect_and_match(flat_gray, persp_gray, method=method)
+
+        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)
 
-        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
-        persp_out = persp_img.copy()
-        cv2.polylines(persp_out, [roi_corners_persp.astype(int)], True, (0,255,0), 3)
+        roi_corners_persp = cv2.perspectiveTransform(roi_corners_flat.reshape(-1, 1, 2), H).reshape(-1, 2)
 
-        # Draw attractive detector label box
-        text = det
-        font = cv2.FONT_HERSHEY_SIMPLEX
-        scale = 1
-        thickness = 2
-        text_size = cv2.getTextSize(text, font, scale, thickness)[0]
-        text_x, text_y = 10, 30
-        cv2.rectangle(persp_out, (text_x-5, text_y-25), (text_x + text_size[0]+5, text_y+5), (0,0,0), -1)
-        cv2.putText(persp_out, text, (text_x, text_y), font, scale, (0,255,255), thickness)
+        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 to RGB for Gradio
-        persp_rgb = cv2.cvtColor(persp_out, cv2.COLOR_BGR2RGB)
-        gallery_images.append(persp_rgb)
+        # Convert BGR to RGB for Gradio display
+        persp_debug_rgb = cv2.cvtColor(persp_debug, cv2.COLOR_BGR2RGB)
+        outputs.append((persp_debug_rgb, method))
 
-    return gallery_images
+    return outputs
 
 # === Gradio Interface ===
-iface = gr.Interface(
-    fn=process_images,
-    inputs=[
-        gr.Image(type="numpy", label="Flat Image"),
-        gr.Image(type="numpy", label="Perspective Image"),
-        gr.File(type="filepath", label="JSON File")
-    ],
-    outputs=[
-        gr.Gallery(label="Detector Results (Perspective Images with ROI & Detector Label)")
-    ],
-    title="Feature Detection ROI Projection",
-    description="Shows SIFT, ORB, BRISK, AKAZE, KAZE detector results on the Perspective image only, with ROI and attractive detector label."
-)
-
-iface.launch()
+with gr.Blocks() as demo:
+    gr.Markdown("## Homography ROI Demo with Multiple Feature Detectors")
+    with gr.Row():
+        flat_input = gr.File(label="Upload Flat Image")
+        persp_input = gr.File(label="Upload Perspective Image")
+        json_input = gr.File(label="Upload mockup.json")
+    output_gallery = gr.Gallery(label="Perspective ROI Results").style(grid=[2], height="400px")
+    run_btn = gr.Button("Run Homography")
+
+    run_btn.click(homography_demo, 
+                  inputs=[flat_input, persp_input, json_input],
+                  outputs=output_gallery)
+
+demo.launch()