app.py

import argparse
import gradio as gr

from hloc import extract_features
from extra_utils.utils import (
    matcher_zoo,
    device,
    match_dense,
    match_features,
    get_model,
    get_feature_model,
    display_matches
)

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

    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):
    return {"value": None, "source": choice, "__type__": "update"}


def ui_reset_state(
    match_threshold, extract_max_keypoints, keypoint_threshold, key, image0, image1
):
    match_threshold = 0.2
    extract_max_keypoints = 1000
    keypoint_threshold = 0.015
    key = list(matcher_zoo.keys())[0]
    image0 = None
    image1 = None
    return (
        match_threshold,
        extract_max_keypoints,
        keypoint_threshold,
        key,
        image0,
        image1,
        {"value": None, "source": "upload", "__type__": "update"},
        {"value": None, "source": "upload", "__type__": "update"},
        "upload",
        None,
        {},
        {},
    )


def run(config):
    with gr.Blocks(
        theme=gr.themes.Monochrome(), css="footer {visibility: hidden}"
    ) as app:
        gr.Markdown(
            """
            <p align="center">
            <h1 align="center">Image Matching WebUI</h1> 
            </p>
            """
        )

        with gr.Row(equal_height=False):
            with gr.Column():
                with gr.Row():
                    matcher_list = gr.Dropdown(
                        choices=list(matcher_zoo.keys()),
                        value="disk+lightglue",
                        label="Matching Model",
                        interactive=True,
                    )
                    match_image_src = gr.Radio(
                        ["upload", "webcam", "canvas"],
                        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",
                        type="numpy",
                        interactive=True,
                        image_mode="RGB",
                    )
                    input_image1 = gr.Image(
                        label="Image 1",
                        type="numpy",
                        interactive=True,
                        image_mode="RGB",
                    )

                with gr.Row():
                    button_reset = gr.Button(label="Reset", value="Reset")
                    button_run = gr.Button(
                        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())}
                        """
                    )

                # collect inputs
                inputs = [
                    match_setting_threshold,
                    match_setting_max_features,
                    detect_keypoints_threshold,
                    matcher_list,
                    input_image0,
                    input_image1,
                ]

                # 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,
                        inputs=inputs,
                        outputs=[],
                        fn=run_matching,
                        cache_examples=False,
                        label="Examples (click one of the images below to Run Match)",
                    )

            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")

            # callbacks
            match_image_src.change(
                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
            )

            # collect outputs
            outputs = [
                output_mkpts,
                matches_result_info,
                matcher_info,
            ]
            # button callbacks
            button_run.click(fn=run_matching, inputs=inputs, outputs=outputs)

            # Reset images
            reset_outputs = [
                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,
            ]
            button_reset.click(fn=ui_reset_state, inputs=inputs, outputs=reset_outputs)
        
    app.launch(share=False)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config_path", type=str, default="config.yaml", help="configuration file path"
    )
    args = parser.parse_args()
    config = None
    run(config)