Update app.py

app.py

@@ -1,65 +1,110 @@
-import cv2
-import numpy as np
-import gradio as gr
-
-def sift_ransac_matching(image, template):
-    # Convert to grayscale
-    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    gray_template = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
-
-    # Initialize SIFT detector
-    sift = cv2.SIFT_create()
-
-    # Find the keypoints and descriptors with SIFT
-    kp1, des1 = sift.detectAndCompute(gray_image, None)
-    kp2, des2 = sift.detectAndCompute(gray_template, None)
-
-    # BFMatcher with default params
-    bf = cv2.BFMatcher()
-    matches = bf.knnMatch(des1, des2, k=2)
-
-    # Apply ratio test
-    good_matches = []
-    for m, n in matches:
-        if m.distance < 0.75 * n.distance:
-            good_matches.append(m)
-
-    if len(good_matches) > 4:
-        # Extract location of good 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)
-
-        # Use RANSAC to find homography
-        M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
-
-        if M is not None:
-            # Compute the match score
-            num_matches = np.sum(mask)
-            match_score = num_matches / len(good_matches)  # Matching score as a ratio
-
-            # Set a threshold for match score
-            threshold = 0.8  # Adjust this threshold as needed
-
-            # Determine if the template is found based on the match score
-            if match_score >= threshold:
-                return "Template found"
-            else:
-                return "Template not found"
-        else:
-            return "Template not found"
-    else:
-        return "Not enough good matches are found. Template not found"
-
-# Gradio interface
-iface = gr.Interface(
-    fn=sift_ransac_matching,
-    inputs=[
-        gr.Image(type="numpy", label="Image"),
-        gr.Image(type="numpy", label="Template"),
-    ],
-    outputs=gr.Text(label="Result"),
-    title="Advanced Template Matching",
-    description="Upload an image and a template to check if the template is present in the image using SIFT and RANSAC.",
-)
-
-iface.launch()
+import cv2
+import numpy as np
+import gradio as gr
+
+def sift_ransac_matching(image, template, ratio_thresh=0.75, ransac_reproj_thresh=5.0, score_thresh=0.8, min_inliers=10):
+    """
+    Returns:
+      - result_text (str)
+      - match_score (float, inlier_ratio = inliers / good_matches)
+      - detected (bool)
+      - annotated_image (numpy array, BGR)
+    """
+    if image is None or template is None:
+        return "Invalid input images.", 0.0, False, image
+
+    # Convert to grayscale
+    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    gray_template = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
+
+    # Initialize SIFT detector
+    sift = cv2.SIFT_create()
+
+    # Find keypoints and descriptors
+    kp_img, des_img = sift.detectAndCompute(gray_image, None)
+    kp_tmpl, des_tmpl = sift.detectAndCompute(gray_template, None)
+
+    if des_img is None or des_tmpl is None or len(kp_img) == 0 or len(kp_tmpl) == 0:
+        return "No features detected. Template not found.", 0.0, False, image
+
+    # KNN matches with Lowe's ratio test
+    bf = cv2.BFMatcher()
+    knn = bf.knnMatch(des_img, des_tmpl, k=2)
+
+    good_matches = []
+    for pair in knn:
+        if len(pair) < 2:
+            continue
+        m, n = pair
+        if m.distance < ratio_thresh * n.distance:
+            good_matches.append(m)
+
+    if len(good_matches) < 4:
+        return f"Not enough good matches ({len(good_matches)}). Template not found.", 0.0, False, image
+
+    # Build point arrays
+    src_pts = np.float32([kp_img[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)   # image -> src
+    dst_pts = np.float32([kp_tmpl[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)  # template -> dst
+
+    # Homography with RANSAC
+    M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, ransac_reproj_thresh)
+    if M is None or mask is None:
+        return "Homography failed. Template not found.", 0.0, False, image
+
+    inliers = int(mask.ravel().sum())
+    match_score = float(inliers) / float(len(good_matches))  # inlier ratio
+
+    # Detection decision (both quality and absolute inlier count can help stability)
+    detected = (match_score >= score_thresh) and (inliers >= min_inliers)
+
+    # Annotate image by projecting template corners back to image space
+    annotated = image.copy()
+    if detected:
+        try:
+            # M maps image -> template; invert to map template -> image
+            Minv = np.linalg.inv(M)
+            h, w = gray_template.shape[:2]
+            tmpl_corners = np.float32([[0, 0], [w, 0], [w, h], [0, h]]).reshape(-1, 1, 2)
+            proj = cv2.perspectiveTransform(tmpl_corners, Minv)
+            proj = proj.astype(int)
+
+            # Draw polygon
+            cv2.polylines(annotated, [proj.reshape(-1, 2)], True, (0, 255, 0), 2)
+            # Put label
+            txt = f"Detected | score={match_score:.3f} | inliers={inliers}/{len(good_matches)}"
+            x, y = proj.reshape(-1, 2)[0]
+            cv2.putText(annotated, txt, (max(0, x), max(20, y)),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+        except np.linalg.LinAlgError:
+            # If inversion fails, still return score/decision
+            pass
+
+    result_text = (
+        f"Template {'found' if detected else 'not found'} | "
+        f"score={match_score:.3f} | inliers={inliers}/{len(good_matches)} good matches"
+    )
+    return result_text, match_score, detected, annotated
+
+# Gradio UI
+iface = gr.Interface(
+    fn=sift_ransac_matching,
+    inputs=[
+        gr.Image(type="numpy", label="Target Image"),
+        gr.Image(type="numpy", label="Template Image"),
+        gr.Slider(0.5, 0.95, value=0.75, step=0.01, label="Lowe ratio threshold"),
+        gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="RANSAC reprojection threshold"),
+        gr.Slider(0.1, 1.0, value=0.8, step=0.01, label="Match score threshold"),
+        gr.Slider(0, 100, value=10, step=1, label="Minimum inliers"),
+    ],
+    outputs=[
+        gr.Text(label="Result"),
+        gr.Number(label="Match score (inlier ratio)"),
+        gr.Checkbox(label="Detected?"),
+        gr.Image(label="Annotated result"),
+    ],
+    title="SIFT + RANSAC Template Search",
+    description="Uploads a target and a template. Shows detection decision, match score, and an annotated result."
+)
+
+if __name__ == "__main__":
+    iface.launch()