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", 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": | |
| # 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:]) | |
| # ) | |
| # fig = plot_images( | |
| # [K.tensor_to_image(torch_img2), K.tensor_to_image(img1_warp_to2)], | |
| # ["Image 2", "Image 1 wrapped to 2"], | |
| # ) | |
| return fig | |
| description = """Line Segment Matching with Kornia | |
| In this space you can try out Line 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", | |
| #"Line Segment Homography Warping", | |
| ], | |
| value="Line Matches", | |
| label="Options"), | |
| ], | |
| outputs=gr.components.Plot(), | |
| examples=[["terrace0.JPG", "terrace1.JPG"]], | |
| title="Line Segment Matching with Kornia", | |
| description=description, | |
| ).launch() |