update: ui

app.py

@@ -1,59 +1,20 @@
 import argparse
 import gradio as gr
-
-from hloc import extract_features
 from common.utils import (
     matcher_zoo,
-    device,
-    match_dense,
-    match_features,
-    get_model,
-    get_feature_model,
-    display_matches,
+    change_estimate_geom,
+    run_matching,
+    ransac_zoo,
+    gen_examples,
 )
 
+DESCRIPTION = """
+# Image Matching WebUI
+This Space demonstrates [Image Matching WebUI](https://github.com/Vincentqyw/image-matching-webui) by vincent qin. Feel free to play with it, or duplicate to run image matching without a queue!
 
-def run_matching(
-    match_threshold, extract_max_keypoints, keypoint_threshold, key, image0, image1
-):
-    # image0 and image1 is RGB mode
-    if image0 is None or image1 is None:
-        raise gr.Error("Error: No images found! Please upload two images.")
-
-    model = matcher_zoo[key]
-    match_conf = model["config"]
-    # update match config
-    match_conf["model"]["match_threshold"] = match_threshold
-    match_conf["model"]["max_keypoints"] = extract_max_keypoints
+🔎 For more details about supported local features and matchers, please refer to https://github.com/Vincentqyw/image-matching-webui
 
-    matcher = get_model(match_conf)
-    if model["dense"]:
-        pred = match_dense.match_images(
-            matcher, image0, image1, match_conf["preprocessing"], device=device
-        )
-        del matcher
-        extract_conf = None
-    else:
-        extract_conf = model["config_feature"]
-        # update extract config
-        extract_conf["model"]["max_keypoints"] = extract_max_keypoints
-        extract_conf["model"]["keypoint_threshold"] = keypoint_threshold
-        extractor = get_feature_model(extract_conf)
-        pred0 = extract_features.extract(
-            extractor, image0, extract_conf["preprocessing"]
-        )
-        pred1 = extract_features.extract(
-            extractor, image1, extract_conf["preprocessing"]
-        )
-        pred = match_features.match_images(matcher, pred0, pred1)
-        del extractor
-    fig, num_inliers = display_matches(pred)
-    del pred
-    return (
-        fig,
-        {"matches number": num_inliers},
-        {"match_conf": match_conf, "extractor_conf": extract_conf},
-    )
+"""
 
 
 def ui_change_imagebox(choice):
@@ -61,7 +22,18 @@ def ui_change_imagebox(choice):
 
 
 def ui_reset_state(
-    match_threshold, extract_max_keypoints, keypoint_threshold, key, image0, image1
+    image0,
+    image1,
+    match_threshold,
+    extract_max_keypoints,
+    keypoint_threshold,
+    key,
+    enable_ransac=False,
+    ransac_method="RANSAC",
+    ransac_reproj_threshold=8,
+    ransac_confidence=0.999,
+    ransac_max_iter=10000,
+    choice_estimate_geom="Homography",
 ):
     match_threshold = 0.2
     extract_max_keypoints = 1000
@@ -69,31 +41,35 @@ def ui_reset_state(
     key = list(matcher_zoo.keys())[0]
     image0 = None
     image1 = None
+    enable_ransac = False
     return (
+        image0,
+        image1,
         match_threshold,
         extract_max_keypoints,
         keypoint_threshold,
         key,
-        image0,
-        image1,
-        {"value": None, "source": "upload", "__type__": "update"},
-        {"value": None, "source": "upload", "__type__": "update"},
+        ui_change_imagebox("upload"),
+        ui_change_imagebox("upload"),
         "upload",
         None,
         {},
         {},
+        None,
+        {},
+        False,
+        "RANSAC",
+        8,
+        0.999,
+        10000,
+        "Homography",
     )
 
 
+# "footer {visibility: hidden}"
 def run(config):
-    with gr.Blocks(css="footer {visibility: hidden}") as app:
-        gr.Markdown(
-            """
-            <p align="center">
-            <h1 align="center">Image Matching WebUI</h1> 
-            </p>
-            """
-        )
+    with gr.Blocks(css="style.css") as app:
+        gr.Markdown(DESCRIPTION)
 
         with gr.Row(equal_height=False):
             with gr.Column():
@@ -109,43 +85,6 @@ def run(config):
                         label="Image Source",
                         value="upload",
                     )
-
-                with gr.Row():
-                    match_setting_threshold = gr.Slider(
-                        minimum=0.0,
-                        maximum=1,
-                        step=0.001,
-                        label="Match threshold",
-                        value=0.1,
-                    )
-                    match_setting_max_features = gr.Slider(
-                        minimum=10,
-                        maximum=10000,
-                        step=10,
-                        label="Max number of features",
-                        value=1000,
-                    )
-                # TODO: add line settings
-                with gr.Row():
-                    detect_keypoints_threshold = gr.Slider(
-                        minimum=0,
-                        maximum=1,
-                        step=0.001,
-                        label="Keypoint threshold",
-                        value=0.015,
-                    )
-                    detect_line_threshold = gr.Slider(
-                        minimum=0.1,
-                        maximum=1,
-                        step=0.01,
-                        label="Line threshold",
-                        value=0.2,
-                    )
-                    # matcher_lists = gr.Radio(
-                    #     ["NN-mutual", "Dual-Softmax"],
-                    #     label="Matcher mode",
-                    #     value="NN-mutual",
-                    # )
                 with gr.Row():
                     input_image0 = gr.Image(
                         label="Image 0",
@@ -166,89 +105,147 @@ def run(config):
                         label="Run Match", value="Run Match", variant="primary"
                     )
 
-                with gr.Accordion("Open for More!", open=False):
-                    gr.Markdown(
-                        f"""
-                        <h3>Supported Algorithms</h3>
-                        {", ".join(matcher_zoo.keys())}
-                        """
-                    )
+                with gr.Accordion("Advanced Setting", open=False):
+                    with gr.Accordion("Matching Setting", open=True):
+                        with gr.Row():
+                            match_setting_threshold = gr.Slider(
+                                minimum=0.0,
+                                maximum=1,
+                                step=0.001,
+                                label="Match thres.",
+                                value=0.1,
+                            )
+                            match_setting_max_features = gr.Slider(
+                                minimum=10,
+                                maximum=10000,
+                                step=10,
+                                label="Max features",
+                                value=1000,
+                            )
+                        # TODO: add line settings
+                        with gr.Row():
+                            detect_keypoints_threshold = gr.Slider(
+                                minimum=0,
+                                maximum=1,
+                                step=0.001,
+                                label="Keypoint thres.",
+                                value=0.015,
+                            )
+                            detect_line_threshold = gr.Slider(
+                                minimum=0.1,
+                                maximum=1,
+                                step=0.01,
+                                label="Line thres.",
+                                value=0.2,
+                            )
+                        # matcher_lists = gr.Radio(
+                        #     ["NN-mutual", "Dual-Softmax"],
+                        #     label="Matcher mode",
+                        #     value="NN-mutual",
+                        # )
+                    with gr.Accordion("RANSAC Setting", open=False):
+                        with gr.Row(equal_height=False):
+                            enable_ransac = gr.Checkbox(label="Enable RANSAC")
+                            ransac_method = gr.Dropdown(
+                                choices=ransac_zoo.keys(),
+                                value="RANSAC",
+                                label="RANSAC Method",
+                                interactive=True,
+                            )
+                        ransac_reproj_threshold = gr.Slider(
+                            minimum=0.0,
+                            maximum=12,
+                            step=0.01,
+                            label="Ransac Reproj threshold",
+                            value=8.0,
+                        )
+                        ransac_confidence = gr.Slider(
+                            minimum=0.0,
+                            maximum=1,
+                            step=0.00001,
+                            label="Ransac Confidence",
+                            value=0.99999,
+                        )
+                        ransac_max_iter = gr.Slider(
+                            minimum=0.0,
+                            maximum=100000,
+                            step=100,
+                            label="Ransac Iterations",
+                            value=10000,
+                        )
+
+                    with gr.Accordion("Geometry Setting", open=True):
+                        with gr.Row(equal_height=False):
+                            # show_geom = gr.Checkbox(label="Show Geometry")
+                            choice_estimate_geom = gr.Radio(
+                                ["Fundamental", "Homography"],
+                                label="Reconstruct Geometry",
+                                value="Homography",
+                            )
 
+                # with gr.Column():
                 # collect inputs
                 inputs = [
+                    input_image0,
+                    input_image1,
                     match_setting_threshold,
                     match_setting_max_features,
                     detect_keypoints_threshold,
                     matcher_list,
-                    input_image0,
-                    input_image1,
+                    enable_ransac,
+                    ransac_method,
+                    ransac_reproj_threshold,
+                    ransac_confidence,
+                    ransac_max_iter,
+                    choice_estimate_geom,
                 ]
 
                 # Add some examples
                 with gr.Row():
-                    examples = [
-                        [
-                            0.1,
-                            2000,
-                            0.015,
-                            "disk+lightglue",
-                            "datasets/sacre_coeur/mapping/71295362_4051449754.jpg",
-                            "datasets/sacre_coeur/mapping/93341989_396310999.jpg",
-                        ],
-                        [
-                            0.1,
-                            2000,
-                            0.015,
-                            "loftr",
-                            "datasets/sacre_coeur/mapping/03903474_1471484089.jpg",
-                            "datasets/sacre_coeur/mapping/02928139_3448003521.jpg",
-                        ],
-                        [
-                            0.1,
-                            2000,
-                            0.015,
-                            "disk",
-                            "datasets/sacre_coeur/mapping/10265353_3838484249.jpg",
-                            "datasets/sacre_coeur/mapping/51091044_3486849416.jpg",
-                        ],
-                        [
-                            0.1,
-                            2000,
-                            0.015,
-                            "topicfm",
-                            "datasets/sacre_coeur/mapping/44120379_8371960244.jpg",
-                            "datasets/sacre_coeur/mapping/93341989_396310999.jpg",
-                        ],
-                        [
-                            0.1,
-                            2000,
-                            0.015,
-                            "superpoint+superglue",
-                            "datasets/sacre_coeur/mapping/17295357_9106075285.jpg",
-                            "datasets/sacre_coeur/mapping/44120379_8371960244.jpg",
-                        ],
-                    ]
                     # Example inputs
                     gr.Examples(
-                        examples=examples,
+                        examples=gen_examples(),
                         inputs=inputs,
                         outputs=[],
                         fn=run_matching,
-                        cache_examples=True,
-                        label="Examples (click one of the images below to Run Match)",
+                        cache_examples=False,
+                        label=(
+                            "Examples (click one of the images below to Run"
+                            " Match)"
+                        ),
+                    )
+                with gr.Accordion("Open for More!", open=False):
+                    gr.Markdown(
+                        f"""
+                        <h3>Supported Algorithms</h3>
+                        {", ".join(matcher_zoo.keys())}
+                        """
                     )
 
             with gr.Column():
-                output_mkpts = gr.Image(label="Keypoints Matching", type="numpy")
-                matches_result_info = gr.JSON(label="Matches Statistics")
-                matcher_info = gr.JSON(label="Match info")
+                output_mkpts = gr.Image(
+                    label="Keypoints Matching", type="numpy"
+                )
+                with gr.Accordion(
+                    "Open for More: Matches Statistics", open=False
+                ):
+                    matches_result_info = gr.JSON(label="Matches Statistics")
+                    matcher_info = gr.JSON(label="Match info")
+
+                output_wrapped = gr.Image(label="Wrapped Pair", type="numpy")
+                with gr.Accordion("Open for More: Geometry info", open=False):
+                    geometry_result = gr.JSON(label="Reconstructed Geometry")
 
             # callbacks
             match_image_src.change(
-                fn=ui_change_imagebox, inputs=match_image_src, outputs=input_image0
+                fn=ui_change_imagebox,
+                inputs=match_image_src,
+                outputs=input_image0,
             )
             match_image_src.change(
-                fn=ui_change_imagebox, inputs=match_image_src, outputs=input_image1
+                fn=ui_change_imagebox,
+                inputs=match_image_src,
+                outputs=input_image1,
             )
 
             # collect outputs
@@ -256,34 +253,61 @@ def run(config):
                 output_mkpts,
                 matches_result_info,
                 matcher_info,
+                geometry_result,
+                output_wrapped,
             ]
             # button callbacks
             button_run.click(fn=run_matching, inputs=inputs, outputs=outputs)
 
             # Reset images
             reset_outputs = [
+                input_image0,
+                input_image1,
                 match_setting_threshold,
                 match_setting_max_features,
                 detect_keypoints_threshold,
                 matcher_list,
                 input_image0,
                 input_image1,
-                input_image0,
-                input_image1,
                 match_image_src,
                 output_mkpts,
                 matches_result_info,
                 matcher_info,
+                output_wrapped,
+                geometry_result,
+                enable_ransac,
+                ransac_method,
+                ransac_reproj_threshold,
+                ransac_confidence,
+                ransac_max_iter,
+                choice_estimate_geom,
             ]
-            button_reset.click(fn=ui_reset_state, inputs=inputs, outputs=reset_outputs)
-    app.queue()
+            button_reset.click(
+                fn=ui_reset_state, inputs=inputs, outputs=reset_outputs
+            )
+
+            # estimate geo
+            choice_estimate_geom.change(
+                fn=change_estimate_geom,
+                inputs=[
+                    input_image0,
+                    input_image1,
+                    geometry_result,
+                    choice_estimate_geom,
+                ],
+                outputs=[output_wrapped, geometry_result],
+            )
+
     app.launch(share=False)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--config_path", type=str, default="config.yaml", help="configuration file path"
+        "--config_path",
+        type=str,
+        default="config.yaml",
+        help="configuration file path",
     )
     args = parser.parse_args()
     config = None

common/utils.py

@@ -1,11 +1,14 @@
-import torch
+import os
+import random
 import numpy as np
+import torch
+from itertools import combinations
 import cv2
+import gradio as gr
 from hloc import matchers, extractors
 from hloc.utils.base_model import dynamic_load
 from hloc import match_dense, match_features, extract_features
-from .plotting import draw_matches, fig2im
-from .visualize_util import plot_images, plot_color_line_matches
+from .viz import draw_matches, fig2im, plot_images, plot_color_line_matches
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -22,6 +25,217 @@ def get_feature_model(conf):
     return model
 
 
+def gen_examples():
+    random.seed(1)
+    example_matchers = [
+        "disk+lightglue",
+        "loftr",
+        "disk",
+        "d2net",
+        "topicfm",
+        "superpoint+superglue",
+        "disk+dualsoftmax",
+        "lanet",
+    ]
+
+    def gen_images_pairs(path: str, count: int = 5):
+        imgs_list = [
+            os.path.join(path, file)
+            for file in os.listdir(path)
+            if file.lower().endswith((".jpg", ".jpeg", ".png"))
+        ]
+        pairs = list(combinations(imgs_list, 2))
+        selected = random.sample(range(len(pairs)), count)
+        return [pairs[i] for i in selected]
+    # image pair path
+    path = "datasets/sacre_coeur/mapping"
+    pairs = gen_images_pairs(path, len(example_matchers))
+    match_setting_threshold = 0.1
+    match_setting_max_features = 2000
+    detect_keypoints_threshold = 0.01
+    enable_ransac = False
+    ransac_method = "RANSAC"
+    ransac_reproj_threshold = 8
+    ransac_confidence = 0.999
+    ransac_max_iter = 10000
+    input_lists = []
+    for pair, mt in zip(pairs, example_matchers):
+        input_lists.append(
+            [
+                pair[0],
+                pair[1],
+                match_setting_threshold,
+                match_setting_max_features,
+                detect_keypoints_threshold,
+                mt,
+                enable_ransac,
+                ransac_method,
+                ransac_reproj_threshold,
+                ransac_confidence,
+                ransac_max_iter,
+            ]
+        )
+    return input_lists
+
+
+def filter_matches(
+    pred,
+    ransac_method="RANSAC",
+    ransac_reproj_threshold=8,
+    ransac_confidence=0.999,
+    ransac_max_iter=10000,
+):
+    mkpts0 = None
+    mkpts1 = None
+    feature_type = None
+    if "keypoints0_orig" in pred.keys() and "keypoints1_orig" in pred.keys():
+        mkpts0 = pred["keypoints0_orig"]
+        mkpts1 = pred["keypoints1_orig"]
+        feature_type = "KEYPOINT"
+    elif (
+        "line_keypoints0_orig" in pred.keys()
+        and "line_keypoints1_orig" in pred.keys()
+    ):
+        mkpts0 = pred["line_keypoints0_orig"]
+        mkpts1 = pred["line_keypoints1_orig"]
+        feature_type = "LINE"
+    else:
+        return pred
+    if mkpts0 is None or mkpts0 is None:
+        return pred
+    if ransac_method not in ransac_zoo.keys():
+        ransac_method = "RANSAC"
+    H, mask = cv2.findHomography(
+        mkpts0,
+        mkpts1,
+        method=ransac_zoo[ransac_method],
+        ransacReprojThreshold=ransac_reproj_threshold,
+        confidence=ransac_confidence,
+        maxIters=ransac_max_iter,
+    )
+    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]
+        elif feature_type == "LINE":
+            pred["line_keypoints0_orig"] = mkpts0[mask]
+            pred["line_keypoints1_orig"] = mkpts1[mask]
+    return pred
+
+
+def compute_geom(
+    pred,
+    ransac_method="RANSAC",
+    ransac_reproj_threshold=8,
+    ransac_confidence=0.999,
+    ransac_max_iter=10000,
+) -> dict:
+    mkpts0 = None
+    mkpts1 = None
+
+    if "keypoints0_orig" in pred.keys() and "keypoints1_orig" in pred.keys():
+        mkpts0 = pred["keypoints0_orig"]
+        mkpts1 = pred["keypoints1_orig"]
+
+    if (
+        "line_keypoints0_orig" in pred.keys()
+        and "line_keypoints1_orig" in pred.keys()
+    ):
+        mkpts0 = pred["line_keypoints0_orig"]
+        mkpts1 = pred["line_keypoints1_orig"]
+
+    if mkpts0 is not None and mkpts1 is not None:
+        if len(mkpts0) < 8:
+            return {}
+        h1, w1, _ = pred["image0_orig"].shape
+        geo_info = {}
+        F, inliers = cv2.findFundamentalMat(
+            mkpts0,
+            mkpts1,
+            method=ransac_zoo[ransac_method],
+            ransacReprojThreshold=ransac_reproj_threshold,
+            confidence=ransac_confidence,
+            maxIters=ransac_max_iter,
+        )
+        geo_info["Fundamental"] = F.tolist()
+        H, _ = cv2.findHomography(
+            mkpts1,
+            mkpts0,
+            method=ransac_zoo[ransac_method],
+            ransacReprojThreshold=ransac_reproj_threshold,
+            confidence=ransac_confidence,
+            maxIters=ransac_max_iter,
+        )
+        geo_info["Homography"] = H.tolist()
+        _, H1, H2 = cv2.stereoRectifyUncalibrated(
+            mkpts0.reshape(-1, 2), mkpts1.reshape(-1, 2), F, imgSize=(w1, h1)
+        )
+        geo_info["H1"] = H1.tolist()
+        geo_info["H2"] = H2.tolist()
+        return geo_info
+    else:
+        return {}
+
+
+def wrap_images(img0, img1, geo_info, geom_type):
+    h1, w1, _ = img0.shape
+    h2, w2, _ = img1.shape
+    result_matrix = None
+    if geo_info is not None and len(geo_info) != 0:
+        rectified_image0 = img0
+        rectified_image1 = None
+        H = np.array(geo_info["Homography"])
+        F = np.array(geo_info["Fundamental"])
+        title = []
+        if geom_type == "Homography":
+            rectified_image1 = cv2.warpPerspective(
+                img1, H, (img0.shape[1] + img1.shape[1], img0.shape[0])
+            )
+            result_matrix = H
+            title = ["Image 0", "Image 1 - warped"]
+        elif geom_type == "Fundamental":
+            H1, H2 = np.array(geo_info["H1"]), np.array(geo_info["H2"])
+            rectified_image0 = cv2.warpPerspective(img0, H1, (w1, h1))
+            rectified_image1 = cv2.warpPerspective(img1, H2, (w2, h2))
+            result_matrix = F
+            title = ["Image 0 - warped", "Image 1 - warped"]
+        else:
+            print("Error: Unknown geometry type")
+        fig = plot_images(
+            [rectified_image0.squeeze(), rectified_image1.squeeze()],
+            title,
+            dpi=300,
+        )
+        dictionary = {
+            "row1": result_matrix[0].tolist(),
+            "row2": result_matrix[1].tolist(),
+            "row3": result_matrix[2].tolist(),
+        }
+        return fig2im(fig), dictionary
+    else:
+        return None, None
+
+
+def change_estimate_geom(input_image0, input_image1, matches_info, choice):
+    if (
+        matches_info is None
+        or len(matches_info) < 1
+        or "geom_info" not in matches_info.keys()
+    ):
+        return None, None
+    geom_info = matches_info["geom_info"]
+    wrapped_images = None
+    if choice != "No":
+        wrapped_images, _ = wrap_images(
+            input_image0, input_image1, geom_info, choice
+        )
+        return wrapped_images, matches_info
+    else:
+        return None, None
+
+
 def display_matches(pred: dict):
     img0 = pred["image0_orig"]
     img1 = pred["image1_orig"]
@@ -42,7 +256,10 @@ def display_matches(pred: dict):
             img1,
             mconf,
             dpi=300,
-            titles=["Image 0 - matched keypoints", "Image 1 - matched keypoints"],
+            titles=[
+                "Image 0 - matched keypoints",
+                "Image 1 - matched keypoints",
+            ],
         )
         fig = fig_mkpts
     if "line0_orig" in pred.keys() and "line1_orig" in pred.keys():
@@ -69,13 +286,107 @@ def display_matches(pred: dict):
             else:
                 mconf = np.ones(len(mkpts0))
             fig_mkpts = draw_matches(mkpts0, mkpts1, img0, img1, mconf, dpi=300)
-            fig_lines = cv2.resize(fig_lines, (fig_mkpts.shape[1], fig_mkpts.shape[0]))
+            fig_lines = cv2.resize(
+                fig_lines, (fig_mkpts.shape[1], fig_mkpts.shape[0])
+            )
             fig = np.concatenate([fig_mkpts, fig_lines], axis=0)
         else:
             fig = fig_lines
     return fig, num_inliers
 
 
+def run_matching(
+    image0,
+    image1,
+    match_threshold,
+    extract_max_keypoints,
+    keypoint_threshold,
+    key,
+    enable_ransac=False,
+    ransac_method="RANSAC",
+    ransac_reproj_threshold=8,
+    ransac_confidence=0.999,
+    ransac_max_iter=10000,
+    choice_estimate_geom="Homography",
+):
+    # image0 and image1 is RGB mode
+    if image0 is None or image1 is None:
+        raise gr.Error("Error: No images found! Please upload two images.")
+
+    model = matcher_zoo[key]
+    match_conf = model["config"]
+    # update match config
+    match_conf["model"]["match_threshold"] = match_threshold
+    match_conf["model"]["max_keypoints"] = extract_max_keypoints
+
+    matcher = get_model(match_conf)
+    if model["dense"]:
+        pred = match_dense.match_images(
+            matcher, image0, image1, match_conf["preprocessing"], device=device
+        )
+        del matcher
+        extract_conf = None
+    else:
+        extract_conf = model["config_feature"]
+        # update extract config
+        extract_conf["model"]["max_keypoints"] = extract_max_keypoints
+        extract_conf["model"]["keypoint_threshold"] = keypoint_threshold
+        extractor = get_feature_model(extract_conf)
+        pred0 = extract_features.extract(
+            extractor, image0, extract_conf["preprocessing"]
+        )
+        pred1 = extract_features.extract(
+            extractor, image1, extract_conf["preprocessing"]
+        )
+        pred = match_features.match_images(matcher, pred0, pred1)
+        del extractor
+
+    if enable_ransac:
+        filter_matches(
+            pred,
+            ransac_method=ransac_method,
+            ransac_reproj_threshold=ransac_reproj_threshold,
+            ransac_confidence=ransac_confidence,
+            ransac_max_iter=ransac_max_iter,
+        )
+
+    fig, num_inliers = display_matches(pred)
+    geom_info = compute_geom(pred)
+    output_wrapped, _ = change_estimate_geom(
+        pred["image0_orig"],
+        pred["image1_orig"],
+        {"geom_info": geom_info},
+        choice_estimate_geom,
+    )
+    del pred
+    return (
+        fig,
+        {"matches number": num_inliers},
+        {
+            "match_conf": match_conf,
+            "extractor_conf": extract_conf,
+        },
+        {
+            "geom_info": geom_info,
+        },
+        output_wrapped,
+        # geometry_result,
+    )
+
+
+# @ref: https://docs.opencv.org/4.x/d0/d74/md__build_4_x-contrib_docs-lin64_opencv_doc_tutorials_calib3d_usac.html
+# AND: https://opencv.org/blog/2021/06/09/evaluating-opencvs-new-ransacs
+ransac_zoo = {
+    "RANSAC": cv2.RANSAC,
+    "USAC_MAGSAC": cv2.USAC_MAGSAC,
+    "USAC_DEFAULT": cv2.USAC_DEFAULT,
+    "USAC_FM_8PTS": cv2.USAC_FM_8PTS,
+    "USAC_PROSAC": cv2.USAC_PROSAC,
+    "USAC_FAST": cv2.USAC_FAST,
+    "USAC_ACCURATE": cv2.USAC_ACCURATE,
+    "USAC_PARALLEL": cv2.USAC_PARALLEL,
+}
+
 # Matchers collections
 matcher_zoo = {
     "gluestick": {"config": match_dense.confs["gluestick"], "dense": True},
@@ -147,11 +458,11 @@ matcher_zoo = {
         "config_feature": extract_features.confs["d2net-ss"],
         "dense": False,
     },
-    # "d2net-ms": {
-    #     "config": match_features.confs["NN-mutual"],
-    #     "config_feature": extract_features.confs["d2net-ms"],
-    #     "dense": False,
-    # },
+    "d2net-ms": {
+        "config": match_features.confs["NN-mutual"],
+        "config_feature": extract_features.confs["d2net-ms"],
+        "dense": False,
+    },
     "alike": {
         "config": match_features.confs["NN-mutual"],
         "config_feature": extract_features.confs["alike"],
@@ -177,6 +488,6 @@ matcher_zoo = {
         "config_feature": extract_features.confs["sift"],
         "dense": False,
     },
-    # "roma": {"config": match_dense.confs["roma"], "dense": True},
-    # "DKMv3": {"config": match_dense.confs["dkm"], "dense": True},
+    "roma": {"config": match_dense.confs["roma"], "dense": True},
+    "DKMv3": {"config": match_dense.confs["dkm"], "dense": True},
 }

common/visualize_util.py

@@ -1,642 +0,0 @@
-""" Organize some frequently used visualization functions. """
-import cv2
-import numpy as np
-import matplotlib
-import matplotlib.pyplot as plt
-import copy
-import seaborn as sns
-
-
-# Plot junctions onto the image (return a separate copy)
-def plot_junctions(input_image, junctions, junc_size=3, color=None):
-    """
-    input_image: can be 0~1 float or 0~255 uint8.
-    junctions: Nx2 or 2xN np array.
-    junc_size: the size of the plotted circles.
-    """
-    # Create image copy
-    image = copy.copy(input_image)
-    # Make sure the image is converted to 255 uint8
-    if image.dtype == np.uint8:
-        pass
-    # A float type image ranging from 0~1
-    elif image.dtype in [np.float32, np.float64, np.float] and image.max() <= 2.0:
-        image = (image * 255.0).astype(np.uint8)
-    # A float type image ranging from 0.~255.
-    elif image.dtype in [np.float32, np.float64, np.float] and image.mean() > 10.0:
-        image = image.astype(np.uint8)
-    else:
-        raise ValueError(
-            "[Error] Unknown image data type. Expect 0~1 float or 0~255 uint8."
-        )
-
-    # Check whether the image is single channel
-    if len(image.shape) == 2 or ((len(image.shape) == 3) and (image.shape[-1] == 1)):
-        # Squeeze to H*W first
-        image = image.squeeze()
-
-        # Stack to channle 3
-        image = np.concatenate([image[..., None] for _ in range(3)], axis=-1)
-
-    # Junction dimensions should be N*2
-    if not len(junctions.shape) == 2:
-        raise ValueError("[Error] junctions should be 2-dim array.")
-
-    # Always convert to N*2
-    if junctions.shape[-1] != 2:
-        if junctions.shape[0] == 2:
-            junctions = junctions.T
-        else:
-            raise ValueError("[Error] At least one of the two dims should be 2.")
-
-    # Round and convert junctions to int (and check the boundary)
-    H, W = image.shape[:2]
-    junctions = (np.round(junctions)).astype(np.int)
-    junctions[junctions < 0] = 0
-    junctions[junctions[:, 0] >= H, 0] = H - 1  # (first dim) max bounded by H-1
-    junctions[junctions[:, 1] >= W, 1] = W - 1  # (second dim) max bounded by W-1
-
-    # Iterate through all the junctions
-    num_junc = junctions.shape[0]
-    if color is None:
-        color = (0, 255.0, 0)
-    for idx in range(num_junc):
-        # Fetch one junction
-        junc = junctions[idx, :]
-        cv2.circle(
-            image, tuple(np.flip(junc)), radius=junc_size, color=color, thickness=3
-        )
-
-    return image
-
-
-# Plot line segements given junctions and line adjecent map
-def plot_line_segments(
-    input_image,
-    junctions,
-    line_map,
-    junc_size=3,
-    color=(0, 255.0, 0),
-    line_width=1,
-    plot_survived_junc=True,
-):
-    """
-    input_image: can be 0~1 float or 0~255 uint8.
-    junctions: Nx2 or 2xN np array.
-    line_map: NxN np array
-    junc_size: the size of the plotted circles.
-    color: color of the line segments (can be string "random")
-    line_width: width of the drawn segments.
-    plot_survived_junc: whether we only plot the survived junctions.
-    """
-    # Create image copy
-    image = copy.copy(input_image)
-    # Make sure the image is converted to 255 uint8
-    if image.dtype == np.uint8:
-        pass
-    # A float type image ranging from 0~1
-    elif image.dtype in [np.float32, np.float64, np.float] and image.max() <= 2.0:
-        image = (image * 255.0).astype(np.uint8)
-    # A float type image ranging from 0.~255.
-    elif image.dtype in [np.float32, np.float64, np.float] and image.mean() > 10.0:
-        image = image.astype(np.uint8)
-    else:
-        raise ValueError(
-            "[Error] Unknown image data type. Expect 0~1 float or 0~255 uint8."
-        )
-
-    # Check whether the image is single channel
-    if len(image.shape) == 2 or ((len(image.shape) == 3) and (image.shape[-1] == 1)):
-        # Squeeze to H*W first
-        image = image.squeeze()
-
-        # Stack to channle 3
-        image = np.concatenate([image[..., None] for _ in range(3)], axis=-1)
-
-    # Junction dimensions should be 2
-    if not len(junctions.shape) == 2:
-        raise ValueError("[Error] junctions should be 2-dim array.")
-
-    # Always convert to N*2
-    if junctions.shape[-1] != 2:
-        if junctions.shape[0] == 2:
-            junctions = junctions.T
-        else:
-            raise ValueError("[Error] At least one of the two dims should be 2.")
-
-    # line_map dimension should be 2
-    if not len(line_map.shape) == 2:
-        raise ValueError("[Error] line_map should be 2-dim array.")
-
-    # Color should be "random" or a list or tuple with length 3
-    if color != "random":
-        if not (isinstance(color, tuple) or isinstance(color, list)):
-            raise ValueError("[Error] color should have type list or tuple.")
-        else:
-            if len(color) != 3:
-                raise ValueError(
-                    "[Error] color should be a list or tuple with length 3."
-                )
-
-    # Make a copy of the line_map
-    line_map_tmp = copy.copy(line_map)
-
-    # Parse line_map back to segment pairs
-    segments = np.zeros([0, 4])
-    for idx in range(junctions.shape[0]):
-        # if no connectivity, just skip it
-        if line_map_tmp[idx, :].sum() == 0:
-            continue
-        # record the line segment
-        else:
-            for idx2 in np.where(line_map_tmp[idx, :] == 1)[0]:
-                p1 = np.flip(junctions[idx, :])  # Convert to xy format
-                p2 = np.flip(junctions[idx2, :])  # Convert to xy format
-                segments = np.concatenate(
-                    (segments, np.array([p1[0], p1[1], p2[0], p2[1]])[None, ...]),
-                    axis=0,
-                )
-
-                # Update line_map
-                line_map_tmp[idx, idx2] = 0
-                line_map_tmp[idx2, idx] = 0
-
-    # Draw segment pairs
-    for idx in range(segments.shape[0]):
-        seg = np.round(segments[idx, :]).astype(np.int)
-        # Decide the color
-        if color != "random":
-            color = tuple(color)
-        else:
-            color = tuple(
-                np.random.rand(
-                    3,
-                )
-            )
-        cv2.line(
-            image, tuple(seg[:2]), tuple(seg[2:]), color=color, thickness=line_width
-        )
-
-    # Also draw the junctions
-    if not plot_survived_junc:
-        num_junc = junctions.shape[0]
-        for idx in range(num_junc):
-            # Fetch one junction
-            junc = junctions[idx, :]
-            cv2.circle(
-                image,
-                tuple(np.flip(junc)),
-                radius=junc_size,
-                color=(0, 255.0, 0),
-                thickness=3,
-            )
-    # Only plot the junctions which are part of a line segment
-    else:
-        for idx in range(segments.shape[0]):
-            seg = np.round(segments[idx, :]).astype(np.int)  # Already in HW format.
-            cv2.circle(
-                image,
-                tuple(seg[:2]),
-                radius=junc_size,
-                color=(0, 255.0, 0),
-                thickness=3,
-            )
-            cv2.circle(
-                image,
-                tuple(seg[2:]),
-                radius=junc_size,
-                color=(0, 255.0, 0),
-                thickness=3,
-            )
-
-    return image
-
-
-# Plot line segments given Nx4 or Nx2x2 line segments
-def plot_line_segments_from_segments(
-    input_image, line_segments, junc_size=3, color=(0, 255.0, 0), line_width=1
-):
-    # Create image copy
-    image = copy.copy(input_image)
-    # Make sure the image is converted to 255 uint8
-    if image.dtype == np.uint8:
-        pass
-    # A float type image ranging from 0~1
-    elif image.dtype in [np.float32, np.float64, np.float] and image.max() <= 2.0:
-        image = (image * 255.0).astype(np.uint8)
-    # A float type image ranging from 0.~255.
-    elif image.dtype in [np.float32, np.float64, np.float] and image.mean() > 10.0:
-        image = image.astype(np.uint8)
-    else:
-        raise ValueError(
-            "[Error] Unknown image data type. Expect 0~1 float or 0~255 uint8."
-        )
-
-    # Check whether the image is single channel
-    if len(image.shape) == 2 or ((len(image.shape) == 3) and (image.shape[-1] == 1)):
-        # Squeeze to H*W first
-        image = image.squeeze()
-
-        # Stack to channle 3
-        image = np.concatenate([image[..., None] for _ in range(3)], axis=-1)
-
-    # Check the if line_segments are in (1) Nx4, or (2) Nx2x2.
-    H, W, _ = image.shape
-    # (1) Nx4 format
-    if len(line_segments.shape) == 2 and line_segments.shape[-1] == 4:
-        # Round to int32
-        line_segments = line_segments.astype(np.int32)
-
-        # Clip H dimension
-        line_segments[:, 0] = np.clip(line_segments[:, 0], a_min=0, a_max=H - 1)
-        line_segments[:, 2] = np.clip(line_segments[:, 2], a_min=0, a_max=H - 1)
-
-        # Clip W dimension
-        line_segments[:, 1] = np.clip(line_segments[:, 1], a_min=0, a_max=W - 1)
-        line_segments[:, 3] = np.clip(line_segments[:, 3], a_min=0, a_max=W - 1)
-
-        # Convert to Nx2x2 format
-        line_segments = np.concatenate(
-            [
-                np.expand_dims(line_segments[:, :2], axis=1),
-                np.expand_dims(line_segments[:, 2:], axis=1),
-            ],
-            axis=1,
-        )
-
-    # (2) Nx2x2 format
-    elif len(line_segments.shape) == 3 and line_segments.shape[-1] == 2:
-        # Round to int32
-        line_segments = line_segments.astype(np.int32)
-
-        # Clip H dimension
-        line_segments[:, :, 0] = np.clip(line_segments[:, :, 0], a_min=0, a_max=H - 1)
-        line_segments[:, :, 1] = np.clip(line_segments[:, :, 1], a_min=0, a_max=W - 1)
-
-    else:
-        raise ValueError(
-            "[Error] line_segments should be either Nx4 or Nx2x2 in HW format."
-        )
-
-    # Draw segment pairs (all segments should be in HW format)
-    image = image.copy()
-    for idx in range(line_segments.shape[0]):
-        seg = np.round(line_segments[idx, :, :]).astype(np.int32)
-        # Decide the color
-        if color != "random":
-            color = tuple(color)
-        else:
-            color = tuple(
-                np.random.rand(
-                    3,
-                )
-            )
-        cv2.line(
-            image,
-            tuple(np.flip(seg[0, :])),
-            tuple(np.flip(seg[1, :])),
-            color=color,
-            thickness=line_width,
-        )
-
-        # Also draw the junctions
-        cv2.circle(
-            image,
-            tuple(np.flip(seg[0, :])),
-            radius=junc_size,
-            color=(0, 255.0, 0),
-            thickness=3,
-        )
-        cv2.circle(
-            image,
-            tuple(np.flip(seg[1, :])),
-            radius=junc_size,
-            color=(0, 255.0, 0),
-            thickness=3,
-        )
-
-    return image
-
-
-# Additional functions to visualize multiple images at the same time,
-# e.g. for line matching
-def plot_images(imgs, titles=None, cmaps="gray", dpi=100, size=5, pad=0.5):
-    """Plot a set of images horizontally.
-    Args:
-        imgs: a list of NumPy or PyTorch images, RGB (H, W, 3) or mono (H, W).
-        titles: a list of strings, as titles for each image.
-        cmaps: colormaps for monochrome images.
-    """
-    n = len(imgs)
-    if not isinstance(cmaps, (list, tuple)):
-        cmaps = [cmaps] * n
-    # figsize = (size*n, size*3/4) if size is not None else None
-    figsize = (size * n, size * 6 / 5) if size is not None else None
-    fig, ax = plt.subplots(1, n, figsize=figsize, dpi=dpi)
-
-    if n == 1:
-        ax = [ax]
-    for i in range(n):
-        ax[i].imshow(imgs[i], cmap=plt.get_cmap(cmaps[i]))
-        ax[i].get_yaxis().set_ticks([])
-        ax[i].get_xaxis().set_ticks([])
-        ax[i].set_axis_off()
-        for spine in ax[i].spines.values():  # remove frame
-            spine.set_visible(False)
-        if titles:
-            ax[i].set_title(titles[i])
-    fig.tight_layout(pad=pad)
-    return fig
-
-
-def plot_keypoints(kpts, colors="lime", ps=4):
-    """Plot keypoints for existing images.
-    Args:
-        kpts: list of ndarrays of size (N, 2).
-        colors: string, or list of list of tuples (one for each keypoints).
-        ps: size of the keypoints as float.
-    """
-    if not isinstance(colors, list):
-        colors = [colors] * len(kpts)
-    axes = plt.gcf().axes
-    for a, k, c in zip(axes, kpts, colors):
-        a.scatter(k[:, 0], k[:, 1], c=c, s=ps, linewidths=0)
-
-
-def plot_matches(kpts0, kpts1, color=None, lw=1.5, ps=4, indices=(0, 1), a=1.0):
-    """Plot matches for a pair of existing images.
-    Args:
-        kpts0, kpts1: corresponding keypoints of size (N, 2).
-        color: color of each match, string or RGB tuple. Random if not given.
-        lw: width of the lines.
-        ps: size of the end points (no endpoint if ps=0)
-        indices: indices of the images to draw the matches on.
-        a: alpha opacity of the match lines.
-    """
-    fig = plt.gcf()
-    ax = fig.axes
-    assert len(ax) > max(indices)
-    ax0, ax1 = ax[indices[0]], ax[indices[1]]
-    fig.canvas.draw()
-
-    assert len(kpts0) == len(kpts1)
-    if color is None:
-        color = matplotlib.cm.hsv(np.random.rand(len(kpts0))).tolist()
-    elif len(color) > 0 and not isinstance(color[0], (tuple, list)):
-        color = [color] * len(kpts0)
-
-    if lw > 0:
-        # transform the points into the figure coordinate system
-        transFigure = fig.transFigure.inverted()
-        fkpts0 = transFigure.transform(ax0.transData.transform(kpts0))
-        fkpts1 = transFigure.transform(ax1.transData.transform(kpts1))
-        fig.lines += [
-            matplotlib.lines.Line2D(
-                (fkpts0[i, 0], fkpts1[i, 0]),
-                (fkpts0[i, 1], fkpts1[i, 1]),
-                zorder=1,
-                transform=fig.transFigure,
-                c=color[i],
-                linewidth=lw,
-                alpha=a,
-            )
-            for i in range(len(kpts0))
-        ]
-
-    # freeze the axes to prevent the transform to change
-    ax0.autoscale(enable=False)
-    ax1.autoscale(enable=False)
-
-    if ps > 0:
-        ax0.scatter(kpts0[:, 0], kpts0[:, 1], c=color, s=ps, zorder=2)
-        ax1.scatter(kpts1[:, 0], kpts1[:, 1], c=color, s=ps, zorder=2)
-
-
-def plot_lines(
-    lines, line_colors="orange", point_colors="cyan", ps=4, lw=2, indices=(0, 1)
-):
-    """Plot lines and endpoints for existing images.
-    Args:
-        lines: list of ndarrays of size (N, 2, 2).
-        colors: string, or list of list of tuples (one for each keypoints).
-        ps: size of the keypoints as float pixels.
-        lw: line width as float pixels.
-        indices: indices of the images to draw the matches on.
-    """
-    if not isinstance(line_colors, list):
-        line_colors = [line_colors] * len(lines)
-    if not isinstance(point_colors, list):
-        point_colors = [point_colors] * len(lines)
-
-    fig = plt.gcf()
-    ax = fig.axes
-    assert len(ax) > max(indices)
-    axes = [ax[i] for i in indices]
-    fig.canvas.draw()
-
-    # Plot the lines and junctions
-    for a, l, lc, pc in zip(axes, lines, line_colors, point_colors):
-        for i in range(len(l)):
-            line = matplotlib.lines.Line2D(
-                (l[i, 0, 0], l[i, 1, 0]),
-                (l[i, 0, 1], l[i, 1, 1]),
-                zorder=1,
-                c=lc,
-                linewidth=lw,
-            )
-            a.add_line(line)
-        pts = l.reshape(-1, 2)
-        a.scatter(pts[:, 0], pts[:, 1], c=pc, s=ps, linewidths=0, zorder=2)
-
-    return fig
-
-
-def plot_line_matches(kpts0, kpts1, color=None, lw=1.5, indices=(0, 1), a=1.0):
-    """Plot matches for a pair of existing images, parametrized by their middle point.
-    Args:
-        kpts0, kpts1: corresponding middle points of the lines of size (N, 2).
-        color: color of each match, string or RGB tuple. Random if not given.
-        lw: width of the lines.
-        indices: indices of the images to draw the matches on.
-        a: alpha opacity of the match lines.
-    """
-    fig = plt.gcf()
-    ax = fig.axes
-    assert len(ax) > max(indices)
-    ax0, ax1 = ax[indices[0]], ax[indices[1]]
-    fig.canvas.draw()
-
-    assert len(kpts0) == len(kpts1)
-    if color is None:
-        color = matplotlib.cm.hsv(np.random.rand(len(kpts0))).tolist()
-    elif len(color) > 0 and not isinstance(color[0], (tuple, list)):
-        color = [color] * len(kpts0)
-
-    if lw > 0:
-        # transform the points into the figure coordinate system
-        transFigure = fig.transFigure.inverted()
-        fkpts0 = transFigure.transform(ax0.transData.transform(kpts0))
-        fkpts1 = transFigure.transform(ax1.transData.transform(kpts1))
-        fig.lines += [
-            matplotlib.lines.Line2D(
-                (fkpts0[i, 0], fkpts1[i, 0]),
-                (fkpts0[i, 1], fkpts1[i, 1]),
-                zorder=1,
-                transform=fig.transFigure,
-                c=color[i],
-                linewidth=lw,
-                alpha=a,
-            )
-            for i in range(len(kpts0))
-        ]
-
-    # freeze the axes to prevent the transform to change
-    ax0.autoscale(enable=False)
-    ax1.autoscale(enable=False)
-
-
-def plot_color_line_matches(lines, correct_matches=None, lw=2, indices=(0, 1)):
-    """Plot line matches for existing images with multiple colors.
-    Args:
-        lines: list of ndarrays of size (N, 2, 2).
-        correct_matches: bool array of size (N,) indicating correct matches.
-        lw: line width as float pixels.
-        indices: indices of the images to draw the matches on.
-    """
-    n_lines = len(lines[0])
-    colors = sns.color_palette("husl", n_colors=n_lines)
-    np.random.shuffle(colors)
-    alphas = np.ones(n_lines)
-    # If correct_matches is not None, display wrong matches with a low alpha
-    if correct_matches is not None:
-        alphas[~np.array(correct_matches)] = 0.2
-
-    fig = plt.gcf()
-    ax = fig.axes
-    assert len(ax) > max(indices)
-    axes = [ax[i] for i in indices]
-    fig.canvas.draw()
-
-    # Plot the lines
-    for a, l in zip(axes, lines):
-        # Transform the points into the figure coordinate system
-        transFigure = fig.transFigure.inverted()
-        endpoint0 = transFigure.transform(a.transData.transform(l[:, 0]))
-        endpoint1 = transFigure.transform(a.transData.transform(l[:, 1]))
-        fig.lines += [
-            matplotlib.lines.Line2D(
-                (endpoint0[i, 0], endpoint1[i, 0]),
-                (endpoint0[i, 1], endpoint1[i, 1]),
-                zorder=1,
-                transform=fig.transFigure,
-                c=colors[i],
-                alpha=alphas[i],
-                linewidth=lw,
-            )
-            for i in range(n_lines)
-        ]
-
-    return fig
-
-
-def plot_color_lines(lines, correct_matches, wrong_matches, lw=2, indices=(0, 1)):
-    """Plot line matches for existing images with multiple colors:
-    green for correct matches, red for wrong ones, and blue for the rest.
-    Args:
-        lines: list of ndarrays of size (N, 2, 2).
-        correct_matches: list of bool arrays of size N with correct matches.
-        wrong_matches: list of bool arrays of size (N,) with correct matches.
-        lw: line width as float pixels.
-        indices: indices of the images to draw the matches on.
-    """
-    # palette = sns.color_palette()
-    palette = sns.color_palette("hls", 8)
-    blue = palette[5]  # palette[0]
-    red = palette[0]  # palette[3]
-    green = palette[2]  # palette[2]
-    colors = [np.array([blue] * len(l)) for l in lines]
-    for i, c in enumerate(colors):
-        c[np.array(correct_matches[i])] = green
-        c[np.array(wrong_matches[i])] = red
-
-    fig = plt.gcf()
-    ax = fig.axes
-    assert len(ax) > max(indices)
-    axes = [ax[i] for i in indices]
-    fig.canvas.draw()
-
-    # Plot the lines
-    for a, l, c in zip(axes, lines, colors):
-        # Transform the points into the figure coordinate system
-        transFigure = fig.transFigure.inverted()
-        endpoint0 = transFigure.transform(a.transData.transform(l[:, 0]))
-        endpoint1 = transFigure.transform(a.transData.transform(l[:, 1]))
-        fig.lines += [
-            matplotlib.lines.Line2D(
-                (endpoint0[i, 0], endpoint1[i, 0]),
-                (endpoint0[i, 1], endpoint1[i, 1]),
-                zorder=1,
-                transform=fig.transFigure,
-                c=c[i],
-                linewidth=lw,
-            )
-            for i in range(len(l))
-        ]
-
-
-def plot_subsegment_matches(lines, subsegments, lw=2, indices=(0, 1)):
-    """Plot line matches for existing images with multiple colors and
-        highlight the actually matched subsegments.
-    Args:
-        lines: list of ndarrays of size (N, 2, 2).
-        subsegments: list of ndarrays of size (N, 2, 2).
-        lw: line width as float pixels.
-        indices: indices of the images to draw the matches on.
-    """
-    n_lines = len(lines[0])
-    colors = sns.cubehelix_palette(
-        start=2, rot=-0.2, dark=0.3, light=0.7, gamma=1.3, hue=1, n_colors=n_lines
-    )
-
-    fig = plt.gcf()
-    ax = fig.axes
-    assert len(ax) > max(indices)
-    axes = [ax[i] for i in indices]
-    fig.canvas.draw()
-
-    # Plot the lines
-    for a, l, ss in zip(axes, lines, subsegments):
-        # Transform the points into the figure coordinate system
-        transFigure = fig.transFigure.inverted()
-
-        # Draw full line
-        endpoint0 = transFigure.transform(a.transData.transform(l[:, 0]))
-        endpoint1 = transFigure.transform(a.transData.transform(l[:, 1]))
-        fig.lines += [
-            matplotlib.lines.Line2D(
-                (endpoint0[i, 0], endpoint1[i, 0]),
-                (endpoint0[i, 1], endpoint1[i, 1]),
-                zorder=1,
-                transform=fig.transFigure,
-                c="red",
-                alpha=0.7,
-                linewidth=lw,
-            )
-            for i in range(n_lines)
-        ]
-
-        # Draw matched subsegment
-        endpoint0 = transFigure.transform(a.transData.transform(ss[:, 0]))
-        endpoint1 = transFigure.transform(a.transData.transform(ss[:, 1]))
-        fig.lines += [
-            matplotlib.lines.Line2D(
-                (endpoint0[i, 0], endpoint1[i, 0]),
-                (endpoint0[i, 1], endpoint1[i, 1]),
-                zorder=1,
-                transform=fig.transFigure,
-                c=colors[i],
-                alpha=1,
-                linewidth=lw,
-            )
-            for i in range(n_lines)
-        ]

common/{plotting.py → viz.py}

@@ -6,6 +6,7 @@ import matplotlib.cm as cm
 from PIL import Image
 import torch.nn.functional as F
 import torch
+import seaborn as sns
 
 
 def _compute_conf_thresh(data):
@@ -19,7 +20,77 @@ def _compute_conf_thresh(data):
     return thr
 
 
-# --- VISUALIZATION --- #
+def plot_images(imgs, titles=None, cmaps="gray", dpi=100, size=5, pad=0.5):
+    """Plot a set of images horizontally.
+    Args:
+        imgs: a list of NumPy or PyTorch images, RGB (H, W, 3) or mono (H, W).
+        titles: a list of strings, as titles for each image.
+        cmaps: colormaps for monochrome images.
+    """
+    n = len(imgs)
+    if not isinstance(cmaps, (list, tuple)):
+        cmaps = [cmaps] * n
+    # figsize = (size*n, size*3/4) if size is not None else None
+    figsize = (size * n, size * 6 / 5) if size is not None else None
+    fig, ax = plt.subplots(1, n, figsize=figsize, dpi=dpi)
+
+    if n == 1:
+        ax = [ax]
+    for i in range(n):
+        ax[i].imshow(imgs[i], cmap=plt.get_cmap(cmaps[i]))
+        ax[i].get_yaxis().set_ticks([])
+        ax[i].get_xaxis().set_ticks([])
+        ax[i].set_axis_off()
+        for spine in ax[i].spines.values():  # remove frame
+            spine.set_visible(False)
+        if titles:
+            ax[i].set_title(titles[i])
+    fig.tight_layout(pad=pad)
+    return fig
+
+
+def plot_color_line_matches(lines, correct_matches=None, lw=2, indices=(0, 1)):
+    """Plot line matches for existing images with multiple colors.
+    Args:
+        lines: list of ndarrays of size (N, 2, 2).
+        correct_matches: bool array of size (N,) indicating correct matches.
+        lw: line width as float pixels.
+        indices: indices of the images to draw the matches on.
+    """
+    n_lines = len(lines[0])
+    colors = sns.color_palette("husl", n_colors=n_lines)
+    np.random.shuffle(colors)
+    alphas = np.ones(n_lines)
+    # If correct_matches is not None, display wrong matches with a low alpha
+    if correct_matches is not None:
+        alphas[~np.array(correct_matches)] = 0.2
+
+    fig = plt.gcf()
+    ax = fig.axes
+    assert len(ax) > max(indices)
+    axes = [ax[i] for i in indices]
+    fig.canvas.draw()
+
+    # Plot the lines
+    for a, l in zip(axes, lines):
+        # Transform the points into the figure coordinate system
+        transFigure = fig.transFigure.inverted()
+        endpoint0 = transFigure.transform(a.transData.transform(l[:, 0]))
+        endpoint1 = transFigure.transform(a.transData.transform(l[:, 1]))
+        fig.lines += [
+            matplotlib.lines.Line2D(
+                (endpoint0[i, 0], endpoint1[i, 0]),
+                (endpoint0[i, 1], endpoint1[i, 1]),
+                zorder=1,
+                transform=fig.transFigure,
+                c=colors[i],
+                alpha=alphas[i],
+                linewidth=lw,
+            )
+            for i in range(n_lines)
+        ]
+
+    return fig
 
 
 def make_matching_figure(
@@ -57,7 +128,7 @@ def make_matching_figure(
         axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c="w", s=5)
 
     # draw matches
-    if mkpts0.shape[0] > 1 and mkpts1.shape[0] > 1:
+    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
         fig.canvas.draw()
         transFigure = fig.transFigure.inverted()
         fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
@@ -105,8 +176,12 @@ def _make_evaluation_figure(data, b_id, alpha="dynamic"):
     b_mask = data["m_bids"] == b_id
     conf_thr = _compute_conf_thresh(data)
 
-    img0 = (data["image0"][b_id][0].cpu().numpy() * 255).round().astype(np.int32)
-    img1 = (data["image1"][b_id][0].cpu().numpy() * 255).round().astype(np.int32)
+    img0 = (
+        (data["image0"][b_id][0].cpu().numpy() * 255).round().astype(np.int32)
+    )
+    img1 = (
+        (data["image1"][b_id][0].cpu().numpy() * 255).round().astype(np.int32)
+    )
     kpts0 = data["mkpts0_f"][b_mask].cpu().numpy()
     kpts1 = data["mkpts1_f"][b_mask].cpu().numpy()
 
@@ -131,8 +206,10 @@ def _make_evaluation_figure(data, b_id, alpha="dynamic"):
 
     text = [
         f"#Matches {len(kpts0)}",
-        f"Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}",
-        f"Recall({conf_thr:.2e}) ({100 * recall:.1f}%): {n_correct}/{n_gt_matches}",
+        f"Precision({conf_thr:.2e}) ({100 * precision:.1f}%):"
+        f" {n_correct}/{len(kpts0)}",
+        f"Recall({conf_thr:.2e}) ({100 * recall:.1f}%):"
+        f" {n_correct}/{n_gt_matches}",
     ]
 
     # make the figure
@@ -188,7 +265,9 @@ def error_colormap(err, thr, alpha=1.0):
     assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
     x = 1 - np.clip(err / (thr * 2), 0, 1)
     return np.clip(
-        np.stack([2 - x * 2, x * 2, np.zeros_like(x), np.ones_like(x) * alpha], -1),
+        np.stack(
+            [2 - x * 2, x * 2, np.zeros_like(x), np.ones_like(x) * alpha], -1
+        ),
         0,
         1,
     )
@@ -200,9 +279,13 @@ np.random.shuffle(color_map)
 
 
 def draw_topics(
-    data, img0, img1, saved_folder="viz_topics", show_n_topics=8, saved_name=None
+    data,
+    img0,
+    img1,
+    saved_folder="viz_topics",
+    show_n_topics=8,
+    saved_name=None,
 ):
-
     topic0, topic1 = data["topic_matrix"]["img0"], data["topic_matrix"]["img1"]
     hw0_c, hw1_c = data["hw0_c"], data["hw1_c"]
     hw0_i, hw1_i = data["hw0_i"], data["hw1_i"]
@@ -237,7 +320,10 @@ def draw_topics(
         dim=-1, keepdim=True
     )  # .float() / (n_topics - 1) #* 255 + 1
     # topic1[~mask1_nonzero] = -1
-    label_img0, label_img1 = torch.zeros_like(topic0) - 1, torch.zeros_like(topic1) - 1
+    label_img0, label_img1 = (
+        torch.zeros_like(topic0) - 1,
+        torch.zeros_like(topic1) - 1,
+    )
     for i, k in enumerate(top_topics):
         label_img0[topic0 == k] = color_map[k]
         label_img1[topic1 == k] = color_map[k]
@@ -312,24 +398,30 @@ def draw_topicfm_demo(
     opencv_display=False,
     opencv_title="",
 ):
-    topic_map0, topic_map1 = draw_topics(data, img0, img1, show_n_topics=show_n_topics)
-
-    mask_tm0, mask_tm1 = np.expand_dims(topic_map0 >= 0, axis=-1), np.expand_dims(
-        topic_map1 >= 0, axis=-1
+    topic_map0, topic_map1 = draw_topics(
+        data, img0, img1, show_n_topics=show_n_topics
     )
 
+    mask_tm0, mask_tm1 = np.expand_dims(
+        topic_map0 >= 0, axis=-1
+    ), np.expand_dims(topic_map1 >= 0, axis=-1)
+
     topic_cm0, topic_cm1 = cm.jet(topic_map0 / 99.0), cm.jet(topic_map1 / 99.0)
-    topic_cm0 = cv2.cvtColor(topic_cm0[..., :3].astype(np.float32), cv2.COLOR_RGB2BGR)
-    topic_cm1 = cv2.cvtColor(topic_cm1[..., :3].astype(np.float32), cv2.COLOR_RGB2BGR)
+    topic_cm0 = cv2.cvtColor(
+        topic_cm0[..., :3].astype(np.float32), cv2.COLOR_RGB2BGR
+    )
+    topic_cm1 = cv2.cvtColor(
+        topic_cm1[..., :3].astype(np.float32), cv2.COLOR_RGB2BGR
+    )
     overlay0 = (mask_tm0 * topic_cm0 + (1 - mask_tm0) * img0).astype(np.float32)
     overlay1 = (mask_tm1 * topic_cm1 + (1 - mask_tm1) * img1).astype(np.float32)
 
     cv2.addWeighted(overlay0, topic_alpha, img0, 1 - topic_alpha, 0, overlay0)
     cv2.addWeighted(overlay1, topic_alpha, img1, 1 - topic_alpha, 0, overlay1)
 
-    overlay0, overlay1 = (overlay0 * 255).astype(np.uint8), (overlay1 * 255).astype(
-        np.uint8
-    )
+    overlay0, overlay1 = (overlay0 * 255).astype(np.uint8), (
+        overlay1 * 255
+    ).astype(np.uint8)
 
     h0, w0 = img0.shape[:2]
     h1, w1 = img1.shape[:2]
@@ -338,7 +430,9 @@ def draw_topicfm_demo(
     out_fig[:h0, :w0] = overlay0
     if h0 >= h1:
         start = (h0 - h1) // 2
-        out_fig[start : (start + h1), (w0 + margin) : (w0 + margin + w1)] = overlay1
+        out_fig[
+            start : (start + h1), (w0 + margin) : (w0 + margin + w1)
+        ] = overlay1
     else:
         start = (h1 - h0) // 2
         out_fig[:h0, (w0 + margin) : (w0 + margin + w1)] = overlay1[
@@ -358,7 +452,8 @@ def draw_topicfm_demo(
             img1[start : start + h0] * 255
         ).astype(np.uint8)
 
-    # draw matching lines, this is inspried from https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/master/models/utils.py
+    # draw matching lines, this is inspried from
+    # https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/master/models/utils.py
     mkpts0, mkpts1 = np.round(mkpts0).astype(int), np.round(mkpts1).astype(int)
     mcolor = (np.array(mcolor[:, [2, 1, 0]]) * 255).astype(int)
 

style.css

@@ -0,0 +1,18 @@
+h1 {
+    text-align: center;
+  }
+  
+  #duplicate-button {
+    margin: auto;
+    color: white;
+    background: #1565c0;
+    border-radius: 100vh;
+  }
+  
+  #component-0 {
+    /* max-width: 900px; */
+    margin: auto;
+    padding-top: 1.5rem;
+  }
+
+  footer {visibility: hidden}