File size: 3,888 Bytes
3ef85e9 f8f3a7d 3ef85e9 f8f3a7d 3ef85e9 a49c8fd 3ef85e9 a49c8fd 3ef85e9 a49c8fd 3ef85e9 a49c8fd 3ef85e9 a49c8fd 3ef85e9 a49c8fd 3ef85e9 a49c8fd 3ef85e9 f8f3a7d 3ef85e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import gradio as gr
import sys, os
import torch
import matplotlib.pylab as plt
from PIL import ImageOps
def pump_matching(img1, img2, trained_with_st=False, scale=300, max_scale=1, max_rot=0, use_gpu=False):
img1 = ImageOps.exif_transpose(img1)
img2 = ImageOps.exif_transpose(img2)
use_singlescale = max_scale==1 and max_rot==0
if use_singlescale: # single
from test_singlescale import Main, arg_parser
else:
from test_multiscale import Main, arg_parser
parser = arg_parser()
args_list = ['--img1','dummy','--img2','dummy','--post-filter', '--desc','PUMP-stytrf' if trained_with_st else 'PUMP','--resize',str(scale)]
if not use_gpu:
args_list += ['--device', 'cpu']
if not use_singlescale:
args_list += ['--max-scale',str(max_scale),'--max-rot',str(max_rot)]
args = parser.parse_args(args_list)
corres = Main().run_from_args_with_images(img1, img2, args)
fig1 = plt.figure(1)
plt.clf()
ax1 = plt.gca()
ax1.imshow(img1)
ax1.axis('off')
plt.tight_layout(pad=0)
fig2 = plt.figure(2)
plt.clf()
ax2 = plt.gca()
ax2.imshow(img2)
ax2.axis('off')
plt.tight_layout(pad=0)
from tools.viz import plot_grid
if corres.shape[-1] > 4:
corres = corres[corres[:,4]>0,:] # select non-null correspondences
if corres.shape[0]>0: plot_grid(corres, ax1, ax2, marker='+')
img1 = None
img2 = None
return fig1, fig2
has_cuda = torch.cuda.is_available() and torch.cuda.device_count()>0
title = "PUMP local descriptor demo"
description = "This is a visualization demo for the PUMP local descriptors presented in our CVPR 2022 paper <b><a href='https://europe.naverlabs.com/research/publications/pump-pyramidal-and-uniqueness-matching-priors-for-unsupervised-learning-of-local-features/' target='_blank'>PUMP: Pyramidal and Uniqueness Matching Priors for Unsupervised Learning of Local Features</a></b>.</p><p><b>WARNING:</b> this demo runs on cpus with downscaled images, without multi-scale or multi-rotations testing, due to limited memory and computational resources, please check out our <a href='https://github.com/naver/pump' target='_blank'>original github repo</a> for these features.</p>"
article = "<p style='text-align: center'><a href='https://github.com/naver/pump' target='_blank'>Original Github Repo</a></p>"
iface = gr.Interface(
fn=pump_matching,
inputs=[
gr.inputs.Image(shape=None, type="pil", label="First Image"),
gr.inputs.Image(shape=None, type="pil", label="Second Image"),
gr.inputs.Checkbox(default=False, label="Use the model trained with style transfer"),
#gr.inputs.Slider(minimum=300, maximum=600, default=400, step=1, label="Original test scale"),
#gr.inputs.Slider(minimum=1, maximum=4, default=1, step=0.1, label="Multi Scale Testing - maximum scale (makes it slower)"),
#gr.inputs.Slider(minimum=0, maximum=180, default=0, step=45, label="Multi Rotation Testing - max rot (makes it slower)"),]
#+ ([gr.inputs.Checkbox(default=True, label='Use GPU instead of CPU')] if has_cuda else []),"""
],
outputs=[
gr.outputs.Image(type="plot", label="Matches in the first image"),
gr.outputs.Image(type="plot", label="Matches in the second image"),
],
title=title,
theme='peach',
description=description,
article=article,
examples=[
['datasets/gradio_demo/cat_src.jpg','datasets/gradio_demo/cat_tgt.jpg',False],#,400,1,0]+([True] if has_cuda else []),
['datasets/gradio_demo/food_src.jpg','datasets/gradio_demo/food_tgt.jpg',False],#,400,1,0]+([True] if has_cuda else []),
['datasets/demo_warp/mountains_src.jpg','datasets/demo_warp/mountains_tgt.jpg',False],#,400,1,0]+([True] if has_cuda else []),
]
)
iface.launch(enable_queue=True) |