Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
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_from_linesegments", 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) | |
elif line_style == "Line Segment Homography Warping": | |
_, _, img1_warp_to2 = get_homography_values( | |
matched_lines1, matched_lines2, torch_img1 | |
) | |
fig = plot_images( | |
[K.tensor_to_image(torch_img2), K.tensor_to_image(img1_warp_to2)], | |
["Image 2", "Image 1 wrapped to 2"], | |
) | |
elif line_style == "Matched Lines for Homography Warping": | |
_, correspondence_mask, _ = get_homography_values( | |
matched_lines1, matched_lines2, torch_img1 | |
) | |
plot_color_line_matches( | |
[matched_lines1[correspondence_mask], matched_lines2[correspondence_mask]], | |
fig, | |
lw=2, | |
) | |
return fig | |
def get_homography_values(matched_lines1, matched_lines2, torch_img1): | |
H_ransac, correspondence_mask = ransac( | |
matched_lines1.flip(dims=(2,)), matched_lines2.flip(dims=(2,)) | |
) | |
img1_warp_to2 = K.geometry.warp_perspective( | |
torch_img1[None], H_ransac[None], (torch_img1.shape[1:]) | |
) | |
return H_ransac, correspondence_mask, img1_warp_to2 | |
description = """In this space you can try out Line Detection and Segment Matching with the Kornia library as seen in [this tutorial](https://kornia.github.io/tutorials/#category=Line%20matching). | |
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", | |
"Line Segment Homography Warping", | |
"Matched Lines for Homography Warping", | |
], | |
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() | |