|
import gradio as gr |
|
import kornia as K |
|
import kornia.feature as KF |
|
import torch |
|
import matplotlib |
|
matplotlib.use('Agg') |
|
import numpy as np |
|
from plot_utils import plot_images, plot_lines, plot_color_line_matches |
|
|
|
sold2 = KF.SOLD2(pretrained=True, config=None) |
|
ransac = K.geometry.RANSAC(model_type="homography", inl_th=3.0) |
|
|
|
|
|
def infer(img1, img2, line_style: str): |
|
torch_img1 = K.image_to_tensor(img1).float() / 255.0 |
|
torch_img2 = K.image_to_tensor(img2).float() / 255.0 |
|
|
|
torch_img1_gray = K.color.rgb_to_grayscale(torch_img1) |
|
torch_img2_gray = K.color.rgb_to_grayscale(torch_img2) |
|
|
|
imgs = torch.stack([torch_img1_gray, torch_img2_gray],) |
|
|
|
with torch.inference_mode(): |
|
outputs = sold2(imgs) |
|
|
|
line_seg1 = outputs["line_segments"][0] |
|
line_seg2 = outputs["line_segments"][1] |
|
desc1 = outputs["dense_desc"][0] |
|
desc2 = outputs["dense_desc"][1] |
|
|
|
with torch.inference_mode(): |
|
matches = sold2.match(line_seg1, line_seg2, desc1[None], desc2[None]) |
|
|
|
valid_matches = matches != -1 |
|
match_indices = matches[valid_matches] |
|
|
|
matched_lines1 = line_seg1[valid_matches] |
|
matched_lines2 = line_seg2[match_indices] |
|
|
|
imgs_to_plot = [K.tensor_to_image(torch_img1), K.tensor_to_image(torch_img2)] |
|
|
|
fig = plot_images(imgs_to_plot, ["Image 1 - detected lines", "Image 2 - detected lines"]) |
|
if line_style == "Line Matches": |
|
lines_to_plot = [line_seg1.numpy(), line_seg2.numpy()] |
|
plot_lines(lines_to_plot, fig, ps=3, lw=2, indices={0, 1}) |
|
elif line_style == "Color Line Matches": |
|
plot_color_line_matches([matched_lines1, matched_lines2], fig, lw=2) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
return fig |
|
|
|
|
|
|
|
description = """In this space you can try out Line Detection and Segment Matching with the Kornia library as seen in [this tutorial](https://kornia-tutorials.readthedocs.io/en/latest/line_detection_and_matching_sold2.html). |
|
|
|
Just upload two images of a scene with different view points, choose an option for output and run the demo. |
|
""" |
|
|
|
|
|
Iface = gr.Interface( |
|
fn=infer, |
|
inputs=[gr.components.Image(), |
|
gr.components.Image(), |
|
gr.components.Dropdown(["Line Matches", |
|
"Color Line Matches", |
|
|
|
], |
|
value="Line Matches", |
|
label="Options"), |
|
], |
|
outputs=gr.components.Plot(), |
|
examples=[["terrace0.JPG", "terrace1.JPG", "Line Matches"]], |
|
title="Line Segment Matching with Kornia", |
|
description=description, |
|
).launch() |