update: rerun ransac

common/app_class.py

@@ -9,6 +9,7 @@ from common.utils import (
     load_config,
     get_matcher_zoo,
     run_matching,
+    run_ransac,
     gen_examples,
     GRADIO_VERSION,
 )
@@ -159,7 +160,9 @@ class ImageMatchingApp:
                                 label="Ransac Iterations",
                                 value=self.cfg["defaults"]["ransac_max_iter"],
                             )
-
+                            button_ransac = gr.Button(
+                                value="Rerun RANSAC", variant="primary"
+                            )
                         with gr.Accordion("Geometry Setting", open=False):
                             with gr.Row(equal_height=False):
                                 choice_geometry_type = gr.Radio(
@@ -171,6 +174,7 @@ class ImageMatchingApp:
                                 )
 
                     # collect inputs
+                    state_cache = gr.State({})
                     inputs = [
                         input_image0,
                         input_image1,
@@ -184,6 +188,7 @@ class ImageMatchingApp:
                         ransac_max_iter,
                         choice_geometry_type,
                         gr.State(self.matcher_zoo),
+                        # state_cache,
                     ]
 
                     # Add some examples
@@ -207,7 +212,8 @@ class ImageMatchingApp:
                 with gr.Column():
                     output_keypoints = gr.Image(label="Keypoints", type="numpy")
                     output_matches_raw = gr.Image(
-                        label="Raw Matches", type="numpy"
+                        label="Raw Matches",
+                        type="numpy",
                     )
                     output_matches_ransac = gr.Image(
                         label="Ransac Matches", type="numpy"
@@ -254,6 +260,7 @@ class ImageMatchingApp:
                     matcher_info,
                     geometry_result,
                     output_wrapped,
+                    state_cache,
                 ]
                 # button callbacks
                 button_run.click(
@@ -288,6 +295,22 @@ class ImageMatchingApp:
                     fn=self.ui_reset_state, inputs=None, outputs=reset_outputs
                 )
 
+                # run ransac button action
+                button_ransac.click(
+                    fn=run_ransac,
+                    inputs=[
+                        ransac_method,
+                        ransac_reproj_threshold,
+                        ransac_confidence,
+                        ransac_max_iter,
+                        state_cache,
+                    ],
+                    outputs=[
+                        output_matches_ransac,
+                        matches_result_info,
+                    ],
+                )
+
                 # estimate geo
                 choice_geometry_type.change(
                     fn=generate_warp_images,

common/utils.py

@@ -265,12 +265,13 @@ def filter_matches(
     mask = np.array(mask.ravel().astype("bool"), dtype="bool")
     if H is not None:
         if feature_type == "KEYPOINT":
-            pred["keypoints0_orig"] = mkpts0[mask]
-            pred["keypoints1_orig"] = mkpts1[mask]
-            pred["mconf"] = pred["mconf"][mask]
+            pred["mkeypoints0_orig"] = mkpts0[mask]
+            pred["mkeypoints1_orig"] = mkpts1[mask]
+            pred["mmconf"] = pred["mconf"][mask]
         elif feature_type == "LINE":
-            pred["line_keypoints0_orig"] = mkpts0[mask]
-            pred["line_keypoints1_orig"] = mkpts1[mask]
+            pred["mline_keypoints0_orig"] = mkpts0[mask]
+            pred["mline_keypoints1_orig"] = mkpts1[mask]
+        pred["H"] = H
     return pred
 
 
@@ -440,6 +441,50 @@ def generate_warp_images(
         return None, None
 
 
+def run_ransac(
+    ransac_method: str = DEFAULT_RANSAC_METHOD,
+    ransac_reproj_threshold: int = DEFAULT_RANSAC_REPROJ_THRESHOLD,
+    ransac_confidence: float = DEFAULT_RANSAC_CONFIDENCE,
+    ransac_max_iter: int = DEFAULT_RANSAC_MAX_ITER,
+    state_cache: Dict[str, Any] = None,
+):
+    t1 = time.time()
+    logger.info(
+        f"Run RANSAC matches using: {ransac_method} with threshold: {ransac_reproj_threshold}"
+    )
+    logger.info(
+        f"Run RANSAC matches using: {ransac_confidence} with iter: {ransac_max_iter}"
+    )
+    # if enable_ransac:
+    filter_matches(
+        state_cache,
+        ransac_method=ransac_method,
+        ransac_reproj_threshold=ransac_reproj_threshold,
+        ransac_confidence=ransac_confidence,
+        ransac_max_iter=ransac_max_iter,
+    )
+    gr.Info(f"RANSAC matches done using: {time.time()-t1:.3f}s")
+    logger.info(f"RANSAC matches done using: {time.time()-t1:.3f}s")
+    t1 = time.time()
+
+    # plot images with ransac matches
+    titles = [
+        "Image 0 - Ransac matched keypoints",
+        "Image 1 - Ransac matched keypoints",
+    ]
+    output_matches_ransac, num_matches_ransac = display_matches(
+        state_cache, titles=titles, tag="KPTS_RANSAC"
+    )
+    gr.Info(f"Display matches done using: {time.time()-t1:.3f}s")
+    logger.info(f"Display matches done using: {time.time()-t1:.3f}s")
+    t1 = time.time()
+    num_matches_raw = state_cache["num_matches_raw"]
+    return output_matches_ransac, {
+        "num_matches_raw": num_matches_raw,
+        "num_matches_ransac": num_matches_ransac,
+    }
+
+
 def run_matching(
     image0: np.ndarray,
     image1: np.ndarray,
@@ -496,7 +541,7 @@ def run_matching(
     output_matches_ransac = None
 
     # super slow!
-    if "roma" in key.lower():
+    if "roma" in key.lower() and device == "cpu":
         gr.Info(
             f"Success! Please be patient and allow for about 2-3 minutes."
             f" Due to CPU inference, {key} is quiet slow."
@@ -592,7 +637,7 @@ def run_matching(
         "Image 1 - Ransac matched keypoints",
     ]
     output_matches_ransac, num_matches_ransac = display_matches(
-        pred, titles=titles
+        pred, titles=titles, tag="KPTS_RANSAC"
     )
     gr.Info(f"Display matches done using: {time.time()-t1:.3f}s")
     logger.info(f"Display matches done using: {time.time()-t1:.3f}s")
@@ -607,17 +652,20 @@ def run_matching(
         choice_geometry_type,
     )
     plt.close("all")
-    del pred
+    # del pred
     logger.info(f"TOTAL time: {time.time()-t0:.3f}s")
     gr.Info(f"In summary, total time: {time.time()-t0:.3f}s")
 
+    state_cache = pred
+    state_cache["num_matches_raw"] = num_matches_raw
+    state_cache["num_matches_ransac"] = num_matches_ransac
     return (
         output_keypoints,
         output_matches_raw,
         output_matches_ransac,
         {
-            "number raw matches": num_matches_raw,
-            "number ransac matches": num_matches_ransac,
+            "num_raw_matches": num_matches_raw,
+            "num_ransac_matches": num_matches_ransac,
         },
         {
             "match_conf": match_conf,
@@ -627,6 +675,7 @@ def run_matching(
             "geom_info": geom_info,
         },
         output_wrapped,
+        state_cache,
     )
 
 

common/viz.py

@@ -156,7 +156,11 @@ def make_matching_figure(
         axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c="w", s=5)
 
     # draw matches
-    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
+    if (
+        mkpts0.shape[0] != 0
+        and mkpts1.shape[0] != 0
+        and mkpts0.shape == mkpts1.shape
+    ):
         fig.canvas.draw()
         transFigure = fig.transFigure.inverted()
         fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
@@ -377,6 +381,7 @@ def display_matches(
     titles: List[str] = [],
     texts: List[str] = [],
     dpi: int = 300,
+    tag: str = "KPTS_RAW",  # KPTS_RAW, KPTS_RANSAC, LINES_RAW, LINES_RANSAC,
 ) -> Tuple[np.ndarray, int]:
     """
     Displays the matches between two images.
@@ -393,11 +398,13 @@ def display_matches(
     img1 = pred["image1_orig"]
 
     num_inliers = 0
+    # draw raw matches
     if (
         "keypoints0_orig" in pred
         and "keypoints1_orig" in pred
         and pred["keypoints0_orig"] is not None
         and pred["keypoints1_orig"] is not None
+        and tag == "KPTS_RAW"
     ):
         mkpts0 = pred["keypoints0_orig"]
         mkpts1 = pred["keypoints1_orig"]
@@ -417,11 +424,38 @@ def display_matches(
             texts=texts,
         )
         fig = fig_mkpts
+    elif (
+        "mkeypoints0_orig" in pred
+        and "mkeypoints1_orig" in pred
+        and pred["mkeypoints0_orig"] is not None
+        and pred["mkeypoints1_orig"] is not None
+        and tag == "KPTS_RANSAC"
+    ):  # draw ransac matches
+        mkpts0 = pred["mkeypoints0_orig"]
+        mkpts1 = pred["mkeypoints1_orig"]
+        num_inliers = len(mkpts0)
+        if "mmconf" in pred:
+            mmconf = pred["mmconf"]
+        else:
+            mmconf = np.ones(len(mkpts0))
+        fig_mkpts = draw_matches_core(
+            mkpts0,
+            mkpts1,
+            img0,
+            img1,
+            mmconf,
+            dpi=dpi,
+            titles=titles,
+            texts=texts,
+        )
+        fig = fig_mkpts
+    # TODO: draw lines
     if (
         "line0_orig" in pred
         and "line1_orig" in pred
         and pred["line0_orig"] is not None
         and pred["line1_orig"] is not None
+        # and (tag == "LINES_RAW" or tag == "LINES_RANSAC")
     ):
         # lines
         mtlines0 = pred["line0_orig"]