Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
| import matplotlib.pyplot as plt | |
| import cv2 | |
| import kornia as K | |
| import kornia.feature as KF | |
| import numpy as np | |
| import torch | |
| from kornia_moons.feature import * | |
| import gradio as gr | |
| def load_torch_image(fname): | |
| img = K.image_to_tensor(cv2.imread(fname), False).float() /255. | |
| img = K.color.bgr_to_rgb(img) | |
| return img | |
| def inference(file1,file2): | |
| fname1 = file1.name | |
| fname2 = file2.name | |
| img1 = load_torch_image(fname1) | |
| img2 = load_torch_image(fname2) | |
| matcher = KF.LoFTR(pretrained='outdoor') | |
| input_dict = {"image0": K.color.rgb_to_grayscale(img1), # LofTR works on grayscale images only | |
| "image1": K.color.rgb_to_grayscale(img2)} | |
| with torch.no_grad(): | |
| correspondences = matcher(input_dict) | |
| mkpts0 = correspondences['keypoints0'].cpu().numpy() | |
| mkpts1 = correspondences['keypoints1'].cpu().numpy() | |
| H, inliers = cv2.findFundamentalMat(mkpts0, mkpts1, cv2.USAC_MAGSAC, 0.5, 0.999, 100000) | |
| inliers = inliers > 0 | |
| fig, ax = plt.subplots() | |
| draw_LAF_matches( | |
| KF.laf_from_center_scale_ori(torch.from_numpy(mkpts0).view(1,-1, 2), | |
| torch.ones(mkpts0.shape[0]).view(1,-1, 1, 1), | |
| torch.ones(mkpts0.shape[0]).view(1,-1, 1)), | |
| KF.laf_from_center_scale_ori(torch.from_numpy(mkpts1).view(1,-1, 2), | |
| torch.ones(mkpts1.shape[0]).view(1,-1, 1, 1), | |
| torch.ones(mkpts1.shape[0]).view(1,-1, 1)), | |
| torch.arange(mkpts0.shape[0]).view(-1,1).repeat(1,2), | |
| K.tensor_to_image(img1), | |
| K.tensor_to_image(img2), | |
| inliers, | |
| draw_dict={'inlier_color': (0.2, 1, 0.2), | |
| 'tentative_color': None, | |
| 'feature_color': (0.2, 0.5, 1), 'vertical': False}, ax=ax) | |
| plt.axis('off') | |
| fig.savefig('example.jpg',dpi=100) | |
| return 'example.jpg' | |
| title = "Kornia-Loftr" | |
| description = "Gradio demo for Kornia-Loftr. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below." | |
| article = "<p style='text-align: center'><a href='https://kornia.readthedocs.io/en/latest/'>Open Source Differentiable Computer Vision Library</a> | <a href='https://github.com/kornia/kornia'>Github Repo</a></p>" | |
| gr.Interface( | |
| inference, | |
| [gr.inputs.Image(type="file", label="Input1"),gr.inputs.Image(type="file", label="Input2")], | |
| gr.outputs.Image(type="file", label="Output"), | |
| title=title, | |
| description=description, | |
| article=article, | |
| enable_queue=True | |
| ).launch(debug=True) |