naver/PUMP

1# Copyright 2022-present NAVER Corp.2# CC BY-NC-SA 4.03# Available only for non-commercial use4 5import gradio as gr6import sys, os7import torch8import matplotlib.pylab as plt9from PIL import ImageOps10 11def pump_matching(img1, img2, trained_with_st=False, scale=300, max_scale=1, max_rot=0, use_gpu=False):12 13    img1 = ImageOps.exif_transpose(img1)14    img2 = ImageOps.exif_transpose(img2)15 16    use_singlescale = max_scale==1 and max_rot==017    if use_singlescale: # single 18        from test_singlescale import Main, arg_parser19    else:20        from test_multiscale import Main, arg_parser21    parser = arg_parser()22 23    args_list = ['--img1','dummy','--img2','dummy','--post-filter', '--desc','PUMP-stytrf' if trained_with_st else 'PUMP','--resize',str(scale)]24    if not use_gpu:25        args_list += ['--device', 'cpu']26    if not use_singlescale:27        args_list += ['--max-scale',str(max_scale),'--max-rot',str(max_rot)]28        29    args = parser.parse_args(args_list)30    31    corres = Main().run_from_args_with_images(img1, img2, args)32        33    fig1 = plt.figure(1)34    plt.clf()35    ax1 = plt.gca()36    ax1.imshow(img1)37    ax1.axis('off')38    plt.tight_layout(pad=0)  39    40    fig2 = plt.figure(2)41    plt.clf()42    ax2 = plt.gca()43    ax2.imshow(img2)44    ax2.axis('off')45    plt.tight_layout(pad=0) 46       47    from tools.viz import plot_grid48    if corres.shape[-1] > 4:49        corres = corres[corres[:,4]>0,:] # select non-null correspondences50    if corres.shape[0]>0: plot_grid(corres, ax1, ax2, marker='+')51 52    img1 = None 53    img2 = None54 55    return fig1, fig256 57has_cuda = torch.cuda.is_available() and torch.cuda.device_count()>058 59title = "PUMP local descriptor demo"60description = "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>" 61 62article = "<p style='text-align: center'><a href='https://github.com/naver/pump' target='_blank'>Original Github Repo</a></p>"63 64iface = gr.Interface(65    fn=pump_matching,66    inputs=[67        gr.inputs.Image(shape=None, type="pil", label="First Image"),68        gr.inputs.Image(shape=None, type="pil", label="Second Image"),69        gr.inputs.Checkbox(default=False, label="Use the model trained with style transfer"),70        #gr.inputs.Slider(minimum=300, maximum=600, default=400, step=1, label="Original test scale"),71        #gr.inputs.Slider(minimum=1, maximum=4, default=1, step=0.1, label="Multi Scale Testing - maximum scale (makes it slower)"),72        #gr.inputs.Slider(minimum=0, maximum=180, default=0, step=45, label="Multi Rotation Testing - max rot (makes it slower)"),]73        #+ ([gr.inputs.Checkbox(default=True, label='Use GPU instead of CPU')] if has_cuda else []),"""74        ],75    outputs=[76        gr.outputs.Image(type="plot", label="Matches in the first image"),77        gr.outputs.Image(type="plot", label="Matches in the second image"),78        ],79    title=title,80    theme='peach',81    description=description,82    article=article,83    examples=[84        ['datasets/gradio_demo/cat_src.jpg','datasets/gradio_demo/cat_tgt.jpg',False],#,400,1,0]+([True] if has_cuda else []),85        ['datasets/gradio_demo/food_src.jpg','datasets/gradio_demo/food_tgt.jpg',False],#,400,1,0]+([True] if has_cuda else []),86        ['datasets/demo_warp/mountains_src.jpg','datasets/demo_warp/mountains_tgt.jpg',False],#,400,1,0]+([True] if has_cuda else []),87    ]88)89iface.launch(enable_queue=True)