update: imcui

imcui/ui/app_class.py

@@ -47,6 +47,7 @@ class ImageMatchingApp:
         self.example_data_root = kwargs.get(
             "example_data_root", Path(__file__).parents[1] / "datasets"
         )
+        # final step
         self.init_interface()
 
     def init_matcher_dropdown(self):
@@ -107,6 +108,8 @@ class ImageMatchingApp:
                         with gr.Row():
                             button_reset = gr.Button(value="Reset")
                             button_run = gr.Button(value="Run Match", variant="primary")
+                        with gr.Row():
+                            button_stop = gr.Button(value="Force Stop", variant="stop")
 
                         with gr.Accordion("Advanced Setting", open=False):
                             with gr.Accordion("Image Setting", open=True):
@@ -326,7 +329,14 @@ class ImageMatchingApp:
                         output_pred,
                     ]
                     # button callbacks
-                    button_run.click(fn=run_matching, inputs=inputs, outputs=outputs)
+                    click_event = button_run.click(
+                        fn=run_matching, inputs=inputs, outputs=outputs
+                    )
+                    # stop button
+                    button_stop.click(
+                        fn=None, inputs=None, outputs=None, cancels=[click_event]
+                    )
+
                     # Reset images
                     reset_outputs = [
                         input_image0,
@@ -519,11 +529,11 @@ class ImageMatchingApp:
             markdown_table = "| Algo. | Conference | Code | Project | Paper |\n"
             markdown_table += "| ----- | ---------- | ---- | ------- | ----- |\n"
 
-            for k, v in cfg.items():
-                if not v["info"]["display"]:
+            for _, v in cfg.items():
+                if not v["info"].get("display", True):
                     continue
-                github_link = get_link(v["info"]["github"])
-                project_link = get_link(v["info"]["project"])
+                github_link = get_link(v["info"].get("github", ""))
+                project_link = get_link(v["info"].get("project", ""))
                 paper_link = get_link(
                     v["info"]["paper"],
                     (
@@ -534,8 +544,8 @@ class ImageMatchingApp:
                 )
 
                 markdown_table += "{}|{}|{}|{}|{}\n".format(
-                    v["info"]["name"],  # display name
-                    v["info"]["source"],
+                    v["info"].get("name", ""),
+                    v["info"].get("source", ""),
                     github_link,
                     project_link,
                     paper_link,
@@ -547,11 +557,11 @@ class ImageMatchingApp:
                     continue
                 data.append(
                     [
-                        v["info"]["name"],
-                        v["info"]["source"],
-                        v["info"]["github"],
-                        v["info"]["paper"],
-                        v["info"]["project"],
+                        v["info"].get("name", ""),
+                        v["info"].get("source", ""),
+                        v["info"].get("github", ""),
+                        v["info"].get("paper", ""),
+                        v["info"].get("project", ""),
                     ]
                 )
             tab = gr.Dataframe(

imcui/ui/utils.py

@@ -153,12 +153,14 @@ def parse_match_config(conf):
         return {
             "matcher": match_dense.confs.get(conf["matcher"]),
             "dense": True,
+            "info": conf.get("info", {}),
         }
     else:
         return {
             "feature": extract_features.confs.get(conf["feature"]),
             "matcher": match_features.confs.get(conf["matcher"]),
             "dense": False,
+            "info": conf.get("info", {}),
         }
 
 
@@ -842,6 +844,32 @@ def run_ransac(
     )
 
 
+def generate_fake_outputs(
+    output_keypoints,
+    output_matches_raw,
+    output_matches_ransac,
+    match_conf,
+    extract_conf,
+    pred,
+):
+    return (
+        output_keypoints,
+        output_matches_raw,
+        output_matches_ransac,
+        {},
+        {
+            "match_conf": match_conf,
+            "extractor_conf": extract_conf,
+        },
+        {
+            "geom_info": pred.get("geom_info", {}),
+        },
+        None,
+        None,
+        None,
+    )
+
+
 def run_matching(
     image0: np.ndarray,
     image1: np.ndarray,
@@ -911,12 +939,6 @@ def run_matching(
     output_matches_raw = None
     output_matches_ransac = None
 
-    # super slow!
-    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."
-        )
     t0 = time.time()
     model = matcher_zoo[key]
     match_conf = model["matcher"]
@@ -924,6 +946,15 @@ def run_matching(
     match_conf["model"]["match_threshold"] = match_threshold
     match_conf["model"]["max_keypoints"] = extract_max_keypoints
     cache_key = "{}_{}".format(key, match_conf["model"]["name"])
+
+    efficiency = model["info"].get("efficiency", "high")
+    if efficiency == "low":
+        gr.Warning(
+            "Matcher {} is time-consuming, please wait for a while".format(
+                model["info"].get("name", "unknown")
+            )
+        )
+
     if use_cached_model:
         # because of the model cache, we need to update the config
         matcher = model_cache.cache_model(cache_key, get_model, match_conf)
@@ -934,6 +965,9 @@ def run_matching(
         matcher = get_model(match_conf)
     logger.info(f"Loading model using: {time.time()-t0:.3f}s")
     t1 = time.time()
+    yield generate_fake_outputs(
+        output_keypoints, output_matches_raw, output_matches_ransac, match_conf, {}, {}
+    )
 
     if model["dense"]:
         if not match_conf["preprocessing"].get("force_resize", False):
@@ -997,6 +1031,14 @@ def run_matching(
         "Image 1 - Keypoints",
     ]
     output_keypoints = display_keypoints(pred, titles=titles)
+    yield generate_fake_outputs(
+        output_keypoints,
+        output_matches_raw,
+        output_matches_ransac,
+        match_conf,
+        extract_conf,
+        pred,
+    )
 
     # plot images with raw matches
     titles = [
@@ -1004,6 +1046,14 @@ def run_matching(
         "Image 1 - Raw matched keypoints",
     ]
     output_matches_raw, num_matches_raw = display_matches(pred, titles=titles)
+    yield generate_fake_outputs(
+        output_keypoints,
+        output_matches_raw,
+        output_matches_ransac,
+        match_conf,
+        extract_conf,
+        pred,
+    )
 
     # if enable_ransac:
     filter_matches(
@@ -1026,9 +1076,17 @@ def run_matching(
     output_matches_ransac, num_matches_ransac = display_matches(
         pred, titles=titles, tag="KPTS_RANSAC"
     )
+    yield generate_fake_outputs(
+        output_keypoints,
+        output_matches_raw,
+        output_matches_ransac,
+        match_conf,
+        extract_conf,
+        pred,
+    )
+
     # 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()
     # plot wrapped images
     output_wrapped, warped_image = generate_warp_images(
@@ -1051,7 +1109,8 @@ def run_matching(
     with open(tmp_state_cache, "wb") as f:
         pickle.dump(state_cache, f)
     logger.info("Dump results done!")
-    return (
+
+    yield (
         output_keypoints,
         output_matches_raw,
         output_matches_ransac,