aveworking/SuperGLue

1import matplotlib.cm as cm2import torch3import gradio as gr4from models.matching import Matching5from models.utils import (make_matching_plot_fast, process_image)6 7torch.set_grad_enabled(False)8 9# Load the SuperPoint and SuperGlue models.10device = 'cuda' if torch.cuda.is_available() else 'cpu'11 12resize = [640, 640]13max_keypoints = 102414keypoint_threshold = 0.00515nms_radius = 416sinkhorn_iterations = 2017match_threshold = 0.218resize_float = False19 20config_indoor = {21    'superpoint': {22        'nms_radius': nms_radius,23        'keypoint_threshold': keypoint_threshold,24        'max_keypoints': max_keypoints25    },26    'superglue': {27        'weights': "indoor",28        'sinkhorn_iterations': sinkhorn_iterations,29        'match_threshold': match_threshold,30    }31}32 33config_outdoor = {34    'superpoint': {35        'nms_radius': nms_radius,36        'keypoint_threshold': keypoint_threshold,37        'max_keypoints': max_keypoints38    },39    'superglue': {40        'weights': "outdoor",41        'sinkhorn_iterations': sinkhorn_iterations,42        'match_threshold': match_threshold,43    }44}45 46matching_indoor = Matching(config_indoor).eval().to(device)47matching_outdoor = Matching(config_outdoor).eval().to(device)48 49def run(input0, input1, superglue):50    if superglue == "indoor":51        matching = matching_indoor52    else:53        matching = matching_outdoor54    55    name0 = 'image1'56    name1 = 'image2'57 58    # If a rotation integer is provided (e.g. from EXIF data), use it:59    rot0, rot1 = 0, 060 61    # Load the image pair.62    image0, inp0, scales0 = process_image(input0, device, resize, rot0, resize_float)63    image1, inp1, scales1 = process_image(input1, device, resize, rot1, resize_float)64 65    if image0 is None or image1 is None:66        print('Problem reading image pair')67        return68 69    # Perform the matching.70    pred = matching({'image0': inp0, 'image1': inp1})71    pred = {k: v[0].detach().numpy() for k, v in pred.items()}72    kpts0, kpts1 = pred['keypoints0'], pred['keypoints1']73    matches, conf = pred['matches0'], pred['matching_scores0']74            75    valid = matches > -176    mkpts0 = kpts0[valid]77    mkpts1 = kpts1[matches[valid]]78    mconf = conf[valid]79 80   81    # Visualize the matches.82    color = cm.jet(mconf)83    text = [84        'SuperGlue',85        'Keypoints: {}:{}'.format(len(kpts0), len(kpts1)),86        '{}'.format(len(mkpts0)),87    ]88 89    if rot0 != 0 or rot1 != 0:90        text.append('Rotation: {}:{}'.format(rot0, rot1))91 92    # Display extra parameter info.93    k_thresh = matching.superpoint.config['keypoint_threshold']94    m_thresh = matching.superglue.config['match_threshold']95    small_text = [96        'Keypoint Threshold: {:.4f}'.format(k_thresh),97        'Match Threshold: {:.2f}'.format(m_thresh),98        'Image Pair: {}:{}'.format(name0, name1),99    ]100 101    output = make_matching_plot_fast(102        image0, image1, kpts0, kpts1, mkpts0, mkpts1, color,103        text, show_keypoints=True, small_text=small_text)104 105    print('Source Image - {}, Destination Image - {}, {}, Match Percentage - {}'.format(name0, name1, text[2], len(mkpts0)/len(kpts0)))106    return output, text[2], str((len(mkpts0)/len(kpts0))*100.0) + '%'107 108if __name__ == '__main__':109 110    glue = gr.Interface(111        fn=run, 112        inputs=[113            gr.Image(label='Input Image'),114            gr.Image(label='Match Image'),115            gr.Radio(choices=["indoor", "outdoor"], value="indoor", type="value", label="SuperGlueType", interactive=True),116        ], 117        outputs=[gr.Image(118            type="pil",119            label="Result"),120            gr.Textbox(label="Keypoints Matched"),121            gr.Textbox(label="Match Percentage")122        ],123        examples=[124            ['./taj-1.jpg', './taj-2.jpg', "outdoor"],125            ['./outdoor-1.JPEG', './outdoor-2.JPEG', "outdoor"]126        ]    127    )128    glue.queue()129    glue.launch()