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)