app.py

import matplotlib.cm as cm
import torch
import gradio as gr
from models.matching import Matching
from models.utils import (make_matching_plot_fast, process_image)

torch.set_grad_enabled(False)

# Load the SuperPoint and SuperGlue models.
device = 'cuda' if torch.cuda.is_available() else 'cpu'

resize = [640, 640]
max_keypoints = 1024
keypoint_threshold = 0.005
nms_radius = 4
sinkhorn_iterations = 20
match_threshold = 0.2
resize_float = False

config_indoor = {
 'superpoint': {
 'nms_radius': nms_radius,
 'keypoint_threshold': keypoint_threshold,
 'max_keypoints': max_keypoints
 },
 'superglue': {
 'weights': "indoor",
 'sinkhorn_iterations': sinkhorn_iterations,
 'match_threshold': match_threshold,
 }
}

config_outdoor = {
 'superpoint': {
 'nms_radius': nms_radius,
 'keypoint_threshold': keypoint_threshold,
 'max_keypoints': max_keypoints
 },
 'superglue': {
 'weights': "outdoor",
 'sinkhorn_iterations': sinkhorn_iterations,
 'match_threshold': match_threshold,
 }
}

matching_indoor = Matching(config_indoor).eval().to(device)
matching_outdoor = Matching(config_outdoor).eval().to(device)

def run(input0, input1, superglue):
 if superglue == "indoor":
 matching = matching_indoor
 else:
 matching = matching_outdoor
 
 name0 = 'image1'
 name1 = 'image2'

 # If a rotation integer is provided (e.g. from EXIF data), use it:
 rot0, rot1 = 0, 0

 # Load the image pair.
 image0, inp0, scales0 = process_image(input0, device, resize, rot0, resize_float)
 image1, inp1, scales1 = process_image(input1, device, resize, rot1, resize_float)

 if image0 is None or image1 is None:
 print('Problem reading image pair')
 return

 # Perform the matching.
 pred = matching({'image0': inp0, 'image1': inp1})
 pred = {k: v[0].detach().numpy() for k, v in pred.items()}
 kpts0, kpts1 = pred['keypoints0'], pred['keypoints1']
 matches, conf = pred['matches0'], pred['matching_scores0']
 
 valid = matches > -1
 mkpts0 = kpts0[valid]
 mkpts1 = kpts1[matches[valid]]
 mconf = conf[valid]

 
 # Visualize the matches.
 color = cm.jet(mconf)
 text = [
 'SuperGlue',
 'Keypoints: {}:{}'.format(len(kpts0), len(kpts1)),
 '{}'.format(len(mkpts0)),
 ]

 if rot0 != 0 or rot1 != 0:
 text.append('Rotation: {}:{}'.format(rot0, rot1))

 # Display extra parameter info.
 k_thresh = matching.superpoint.config['keypoint_threshold']
 m_thresh = matching.superglue.config['match_threshold']
 small_text = [
 'Keypoint Threshold: {:.4f}'.format(k_thresh),
 'Match Threshold: {:.2f}'.format(m_thresh),
 'Image Pair: {}:{}'.format(name0, name1),
 ]

 output = make_matching_plot_fast(
 image0, image1, kpts0, kpts1, mkpts0, mkpts1, color,
 text, show_keypoints=True, small_text=small_text)

 print('Source Image - {}, Destination Image - {}, {}, Match Percentage - {}'.format(name0, name1, text[2], len(mkpts0)/len(kpts0)))
 return output, text[2], str((len(mkpts0)/len(kpts0))*100.0) + '%'

if __name__ == '__main__':

 glue = gr.Interface(
 fn=run, 
 inputs=[
 gr.Image(label='Input Image'),
 gr.Image(label='Match Image'),
 gr.Radio(choices=["indoor", "outdoor"], value="indoor", type="value", label="SuperGlueType", interactive=True),
 ], 
 outputs=[gr.Image(
 type="pil",
 label="Result"),
 gr.Textbox(label="Keypoints Matched"),
 gr.Textbox(label="Match Percentage")
 ],
 examples=[
 ['./taj-1.jpg', './taj-2.jpg', "outdoor"],
 ['./outdoor-1.JPEG', './outdoor-2.JPEG', "outdoor"]
 ] 
 )
 glue.queue()
 glue.launch()