edits

app.py

@@ -27,9 +27,13 @@ device = 'cpu'
 # Load net
 state = torch.load('fire.pth', map_location='cpu')
 state['net_params']['pretrained'] = None # no need for imagenet pretrained model
-net = fire_network.init_network(**state['net_params']).to(device)
-net.load_state_dict(state['state_dict'])
+net_sfm = fire_network.init_network(**state['net_params']).to(device)
+net_sfm.load_state_dict(state['state_dict'])
 
+state2 = torch.load('fire_imagenet.pth', map_location='cpu')
+state2['net_params']['pretrained'] = None # no need for imagenet pretrained model
+net_imagenet = fire_network.init_network(**state2['net_params']).to(device)
+net_imagenet.load_state_dict(state2['state_dict'])
 
 # ---------------------------------------
 transform = transforms.Compose([
@@ -86,10 +90,13 @@ def match(query_feat, pos_feat, LoweRatioTh=0.9):
 
 col = plt.get_cmap('tab10')
 
-def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50, sf_ids='', only_matching=True):
+def generate_matching_superfeatures(im1, im2, model_fn, scale_id=6, threshold=50, sf_ids='', only_matching=True):
     print('im1:', im1.size)
     print('im2:', im2.size)
 
+    net = net_sfm
+    if model_fn == "ImageNet":
+        net = net_imagenet
 
     # dataset_ = ImgDataset(images=[im1, im2], imsize=1024)
     # loader = torch.utils.data.DataLoader(dataset_, shuffle=False, pin_memory=True)
@@ -256,8 +263,9 @@ iface = gr.Interface(
 #        gr.inputs.Image(shape=(1024, 1024), type="pil", label="Second Image"),
         gr.inputs.Image(type="pil", label="First Image"),
         gr.inputs.Image(type="pil", label="Second Image"),
+        gr.inputs.Radio(["ImageNet", "SfM-120k (landmarks)"], label="Model", optional=False),
         gr.inputs.Slider(minimum=0, maximum=6, step=1, default=2, label="Scale"),
-        gr.inputs.Slider(minimum=0, maximum=255, step=25, default=100, label="Binarization Threshold"),
+        gr.inputs.Slider(minimum=0, maximum=255, step=25, default=150, label="Binarization Threshold"),
         gr.inputs.Textbox(lines=1, default="", label="SF IDs to show (comma separated numbers from 0-255; typing 'rX' will return X random SFs", optional=True),
         gr.inputs.Checkbox(default=True, label="Show only matching SFs", optional=False),
         ],
@@ -273,9 +281,9 @@ iface = gr.Interface(
     article=article,
     css=css,
     examples=[
-        ["chateau_1.png", "chateau_2.png", 2, 100, '55,14,5,4,52,57,40,9', True],
-        ["anafi1.jpeg", "anafi2.jpeg", 4, 150, '51,141,185,99,', True],
-        ["areopoli1.jpeg", "areopoli2.jpeg", 4, 50, '72,44,142,213,236', True],
+        ["chateau_1.png", "chateau_2.png", 3, 150, 'r8', True],
+        ["anafi1.jpeg", "anafi2.jpeg", 4, 150, 'r8', True],
+        ["areopoli1.jpeg", "areopoli2.jpeg", 4, 150, 'r8', True],
     ]
 )
 iface.launch(enable_queue=True)

fire_imagenet.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e05eb75bd4155cd33d3d6fc063fb658b412fff62b73484f86842fc5df5f17b52
+size 60557200