Update app.py

app.py

@@ -1,178 +1,120 @@
 import gradio as gr
 import cv2
 import numpy as np
-import matplotlib.pyplot as plt
 import json
 import math
-import os
-
-def ransac(image1, image2, detector_type):
-    """
-    Finds the homography matrix using the RANSAC algorithm with the selected feature detector.
-    """
-    gray1 = cv2.cvtColor(image1, cv2.COLOR_RGB2GRAY)
-    gray2 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY)
+import matplotlib.pyplot as plt
 
+# === 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, 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)
+
+def preprocess_gray_clahe(img):
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    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()
-        matcher = cv2.FlannBasedMatcher(dict(algorithm=1, trees=5), dict(checks=50))
-    elif detector_type == "ORB":
-        detector = cv2.ORB_create()
-        matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
+        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, crossCheck=True)
+        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, crossCheck=True)
+        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
     elif detector_type == "KAZE":
         detector = cv2.KAZE_create()
-        matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
-    else:
-        return None
-
-    kp1, des1 = detector.detectAndCompute(gray1, None)
-    kp2, des2 = detector.detectAndCompute(gray2, None)
-
-    if des1 is None or des2 is None or len(kp1) < 2 or len(kp2) < 2:
-        return None
-
-    try:
-        if detector_type == "SIFT":
-            matches = matcher.knnMatch(des1, des2, k=2)
-            good_matches = []
-            if matches:
-                for m, n in matches:
-                    if m.distance < 0.75 * n.distance:
-                        good_matches.append(m)
-        else:
-            matches = matcher.match(des1, des2)
-            good_matches = sorted(matches, key=lambda x: x.distance)
-    except cv2.error as e:
-        print(f"Error during matching: {e}")
-        return None
-
-    if len(good_matches) > 10:
-        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)
-        return H
+        matcher = cv2.BFMatcher(cv2.NORM_L2)
     else:
-        return None
-
-def get_bounding_box_points(json_data):
-    """
-    Extracts and calculates the four corner points of the bounding box, assuming x,y are top-left.
-    """
-    print_area = json_data['printAreas'][0]
-    x = print_area['position']['x']
-    y = print_area['position']['y']
-    w = print_area['width']
-    h = print_area['height']
-    rotation_deg = print_area['rotation']
-
-    points = np.float32([
-        [0, 0],
-        [w, 0],
-        [w, h],
-        [0, h]
-    ]).reshape(-1, 1, 2)
-
-    rotation_rad = math.radians(rotation_deg)
-    cos_theta = math.cos(rotation_rad)
-    sin_theta = math.sin(rotation_rad)
-    rotation_matrix = np.array([
-        [cos_theta, -sin_theta],
-        [sin_theta,  cos_theta]
-    ])
-
-    rotated_points = np.dot(points.reshape(-1, 2), rotation_matrix.T)
-    final_points = rotated_points + np.array([x, y])
-    return final_points.reshape(-1, 1, 2)
-
-def process_and_plot_all_detectors(image1_np, image2_np, json_file):
-    """
-    Processes the images with all available detectors and returns image data for display and download.
-    Keeps original RGB colors intact.
-    """
-    if image1_np is None or image2_np is None:
-        return [None] * 6
-
-    try:
-        with open(json_file.name, 'r') as f:
-            data = json.load(f)
-    except Exception as e:
-        print(f"Error: Could not read JSON file. {e}")
-        return [None] * 6
-
-    detectors = ["SIFT", "ORB", "BRISK", "AKAZE", "KAZE"]
+        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 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
+
+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
+
+def process_images(flat_img, persp_img, json_file):
+    # Preprocess
+    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 = []
-    download_files = [None] * 5
-
-    for i, detector_type in enumerate(detectors):
-        H = ransac(image1_np, image2_np, detector_type)
-
-        if H is not None:
-            box_points = get_bounding_box_points(data)
-            
-            # Convert RGB → BGR for OpenCV drawing
-            output_flat_img = cv2.cvtColor(image1_np, cv2.COLOR_RGB2BGR)
-            cv2.polylines(output_flat_img, [np.int32(box_points)], isClosed=True, color=(0, 0, 255), thickness=5)
-
-            transformed_box_points = cv2.perspectiveTransform(box_points, H)
-
-            output_perspective_img = cv2.cvtColor(image2_np, cv2.COLOR_RGB2BGR)
-            cv2.polylines(output_perspective_img, [np.int32(transformed_box_points)], isClosed=True, color=(0, 0, 255), thickness=5)
-
-            # Convert BGR → RGB for display
-            output_flat_img = cv2.cvtColor(output_flat_img, cv2.COLOR_BGR2RGB)
-            output_perspective_img = cv2.cvtColor(output_perspective_img, cv2.COLOR_BGR2RGB)
-
-            # Plot images side by side
-            fig, axes = plt.subplots(1, 3, figsize=(18, 6))
-            axes[0].imshow(output_flat_img)
-            axes[0].set_title(f'Original (Flat) - {detector_type}')
-            axes[0].axis('off')
-
-            axes[1].imshow(image2_np)  # original perspective image in RGB
-            axes[1].set_title('Original (Perspective)')
-            axes[1].axis('off')
-
-            axes[2].imshow(output_perspective_img)
-            axes[2].set_title('Projected Bounding Box')
-            axes[2].axis('off')
-            
-            plt.tight_layout()
-            file_name = f"result_{detector_type.lower()}.png"
-            plt.savefig(file_name)
-            plt.close(fig)
-
-            gallery_images.append(file_name)
-            download_files[i] = file_name
-        else:
-            print(f"Warning: Homography matrix could not be found with {detector_type} detector. Skipping this result.")
-            # download_files[i] remains None
-
-    return [gallery_images] + download_files
 
+    for det in detectors:
+        kp1, kp2, matches = detect_and_match(flat_gray, persp_gray, det)
+        if len(matches) < 4:
+            # Skip if too few matches
+            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)
+
+    return gallery_images
 
 iface = gr.Interface(
-    fn=process_and_plot_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 AKAZE Result"),
-        gr.File(label="Download KAZE Result")
-    ],
-    title="Homography and Bounding Box Projection with All Detectors",
-    description="Upload two images and a JSON file to see the bounding box projection for all 5 feature extraction methods. Each result can be downloaded separately."
+    fn=process_images,
+    inputs=[gr.Image(type="numpy", label="Flat Image"),
+            gr.Image(type="numpy", label="Perspective Image"),
+            gr.File(label="JSON File")],
+    outputs=gr.Gallery(label="ROI Projection Results"),
+    title="ROI Projection with Multiple Feature Detectors",
+    description="Displays ROI projected from Flat to Perspective image using SIFT, BRISK, ORB, AKAZE, KAZE."
 )
 
-iface.launch()
+iface.launch()