add: force stop button (#113)

* update: ci skip warnings * add: efficiency attribute to config * add: force stop button * update: whl version * update: README
Vincentqyw · Jan 5, 2025 · 7a5cc96 · 7a5cc96
1 parent e527dfe
commit 7a5cc96
Show file tree

Hide file tree

Showing 5 changed files with 144 additions and 30 deletions.
diff --git a/README.md b/README.md
@@ -4,19 +4,25 @@
 [![Issues][issues-shield]][issues-url]
 
 <p align="center">
-  <h1 align="center"><br><ins>$\color{red}{\textnormal{Image\ Matching\ WebUI}}$
-</ins><br>Matching Keypoints between two images</h1>
+  <h1 align="center"><br><ins>Image Matching WebUI</ins>
+  <br>Matching Keypoints between two images</h1>
 </p>
+<div align="center">
+  <a target="_blank" href="https://github.com/Vincentqyw/image-matching-webui/actions/workflows/release.yml"><img src="https://github.com/Vincentqyw/image-matching-webui/actions/workflows/release.yml/badge.svg" alt="PyPI Release"></a>
+  <a target="_blank" href='https://huggingface.co/spaces/Realcat/image-matching-webui'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue'></a>
+  <a target="_blank" href="https://pypi.org/project/imcui"><img alt="PyPI - Version" src="https://img.shields.io/pypi/v/imcui?style=flat&logo=pypi&label=imcui&link=https%3A%2F%2Fpypi.org%2Fproject%2Fimcui"></a>
+  <a target="_blank" href="https://hub.docker.com/r/vincentqin/image-matching-webui"><img alt="Docker Image Version" src="https://img.shields.io/docker/v/vincentqin/image-matching-webui?sort=date&arch=amd64&logo=docker&label=imcui&link=https%3A%2F%2Fhub.docker.com%2Fr%2Fvincentqin%2Fimage-matching-webui"></a>
+
+</div>
 
 ## Description
 
 This simple tool efficiently matches image pairs using multiple famous image matching algorithms. The tool features a Graphical User Interface (GUI) designed using [gradio](https://gradio.app/). You can effortlessly select two images and a matching algorithm and obtain a precise matching result.
 **Note**: the images source can be either local images or webcam images.
 
-Try it on <a href='https://huggingface.co/spaces/Realcat/image-matching-webui'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue'></a>
-<a target="_blank" href="https://lightning.ai/realcat/studios/image-matching-webui">
-  <img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/>
-</a>
+Try it on
+<a href='https://huggingface.co/spaces/Realcat/image-matching-webui'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue'></a>
+<a target="_blank" href="https://lightning.ai/realcat/studios/image-matching-webui"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a>
 
 Here is a demo of the tool:
 

diff --git a/config/config.yaml b/config/config.yaml
@@ -1,6 +1,6 @@
 server:
   name: "0.0.0.0"
-  port: 7861
+  port: 7860
 
 defaults:
   setting_threshold: 0.1
@@ -9,21 +9,48 @@ defaults:
   enable_ransac: true
   ransac_method: CV2_USAC_MAGSAC
   ransac_reproj_threshold: 8
-  ransac_confidence: 0.999
+  ransac_confidence: 0.9999
   ransac_max_iter: 10000
   ransac_num_samples: 4
   match_threshold: 0.2
   setting_geometry: Homography
 
 matcher_zoo:
+  # example config
+  Example:
+    # show in `Matching Model` or not, default: true
+    enable: false
+    # matcher name
+    matcher: example
+    # skip ci or not, default: false
+    skip_ci: true
+    # dense matcher or not, default: true
+    dense: true
+    # info
+    info:
+      # dispaly name in `Matching Model`
+      name: example(example)
+      # conference/journal/workshop Year
+      source: "CVPR XXXX"
+      # github link
+      github: https://github.com/example/example
+      # paper link
+      paper: https://arxiv.org/abs/xxxx.xxxx
+      # project link
+      project: https://example.com
+      # show in `support algos` table
+      display: false
+      # low, medium, high
+      efficiency: low
+
   minima(loftr):
     matcher: minima_loftr
     dense: true
     info:
       name: MINIMA(LoFTR) #dispaly name
       source: "ARXIV 2024"
       paper: https://arxiv.org/abs/2412.19412
-      display: false
+      display: true
   minima(RoMa):
     matcher: minima_roma
     skip_ci: true
@@ -33,6 +60,7 @@ matcher_zoo:
       source: "ARXIV 2024"
       paper: https://arxiv.org/abs/2412.19412
       display: false
+      efficiency: low  # low, medium, high
   omniglue:
     enable: true
     matcher: omniglue
@@ -55,6 +83,7 @@ matcher_zoo:
       paper: https://arxiv.org/abs/2406.09756
       project: https://dust3r.europe.naverlabs.com
       display: true
+      efficiency: low  # low, medium, high
   DUSt3R:
     # TODO: duster is under development
     enable: true
@@ -80,6 +109,7 @@ matcher_zoo:
       paper: https://arxiv.org/abs/2402.11095
       project: https://xuelunshen.com/gim
       display: true
+      efficiency: low  # low, medium, high
   RoMa:
     matcher: roma
     skip_ci: true
@@ -91,6 +121,7 @@ matcher_zoo:
       paper: https://arxiv.org/abs/2305.15404
       project: https://parskatt.github.io/RoMa
       display: true
+      efficiency: low  # low, medium, high
   dkm:
     matcher: dkm
     skip_ci: true
@@ -102,6 +133,7 @@ matcher_zoo:
       paper: https://arxiv.org/abs/2202.00667
       project: https://parskatt.github.io/DKM
       display: true
+      efficiency: low  # low, medium, high
   loftr:
     matcher: loftr
     dense: true
@@ -144,6 +176,7 @@ matcher_zoo:
       paper: https://arxiv.org/abs/2103.14167
       project: null
       display: true
+      efficiency: low  # low, medium, high
   topicfm:
     matcher: topicfm
     dense: true

diff --git a/imcui/ui/app_class.py b/imcui/ui/app_class.py
@@ -47,6 +47,7 @@ def __init__(self, server_name="0.0.0.0", server_port=7860, **kwargs):
         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 @@ def init_interface(self):
                         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 @@ def init_interface(self):
                         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 @@ def get_link(link, tag="Link"):
             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 @@ def get_link(link, tag="Link"):
                 )
 
                 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 @@ def get_link(link, tag="Link"):
                     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(

diff --git a/imcui/ui/utils.py b/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,19 +939,22 @@ 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"]
     # update match config
     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,13 +1031,29 @@ 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 = [
         "Image 0 - Raw matched keypoints",
         "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,

diff --git a/pyproject.toml b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "imcui"
 description = "Image Matching Webui: A tool for matching images using sota algorithms with a Gradio UI"
-version = "0.0.1"
+version = "0.0.2"
 authors = [
     {name = "vincentqyw"},
 ]
@@ -39,3 +39,9 @@ minversion = "6.0"
 addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
 xfail_strict = true
 testpaths = ["tests"]
+filterwarnings = [
+    "ignore::DeprecationWarning",
+    "ignore::UserWarning",
+    "ignore::FutureWarning",
+    "ignore::RuntimeWarning",
+]