edits

app.py

@@ -3,20 +3,18 @@ import gradio as gr
 import cv2
 
 import torch
+import torch.utils.data as data
+from torchvision import transforms
 
 import matplotlib.pyplot as plt
 from matplotlib import cm
 from matplotlib import colors
 from mpl_toolkits.axes_grid1 import ImageGrid
 
-from torchvision import transforms
-
 import fire_network
 
 import numpy as np
 
-
-
 from PIL import Image
 
 # Possible Scales for multiscale inference
@@ -30,11 +28,34 @@ 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'])
 
-transform = transforms.Compose([
-        transforms.Resize(1024),
-        transforms.ToTensor(), 
-        transforms.Normalize(**dict(zip(["mean", "std"], net.runtime['mean_std'])))
-        ])
+
+# ---------------------------------------
+# transform = transforms.Compose([
+#         transforms.Resize(1024),
+#         transforms.ToTensor(), 
+#         transforms.Normalize(**dict(zip(["mean", "std"], net.runtime['mean_std'])))
+#         ])
+# ---------------------------------------
+
+class ImgDataset(data.Dataset):
+    
+    def __init__(self, images, imsize):
+        self.images = images
+        self.imsize = imsize
+        self.transform = transforms.Compose([transforms.ToTensor(), \
+            transforms.Normalize(**dict(zip(["mean", "std"], net.runtime['mean_std'])))])
+
+
+    def __getitem__(self, index):
+        img = self.images[index]
+        return  self.transform(img.thumbnail((self.imsize, self.imsize), Image.ANTIALIAS))
+
+
+    def __len__(self):
+        return len(self.images)
+
+    
+
 
 
 # sf_idx_ = [55, 14, 5, 4, 52, 57, 40, 9]
@@ -51,25 +72,38 @@ def generate_matching_superfeatures(im1, im2, scale_id=6, threshold=50, sf_ids='
         sf_idx_ = np.random.randint(256, size=n_sf_ids)
     elif sf_ids != '':
         sf_idx_ = map(int, sf_ids.strip().split(','))
-    
-    im1_tensor = transform(im1).unsqueeze(0)
-    im2_tensor = transform(im2).unsqueeze(0)
+        
+
+    dataset_ = ImgDataset(images=[im1, im2], imsize=1024)
+    loader = torch.utils.data.DataLoader(dataset_, shuffle=False, pin_memory=True)
+
+
+    # im1_tensor = transform(im1).unsqueeze(0)
+    # im2_tensor = transform(im2).unsqueeze(0)
 
     im1_cv = np.array(im1)[:, :, ::-1].copy() 
     im2_cv = np.array(im2)[:, :, ::-1].copy() 
 
     # extract features
     with torch.no_grad():
-        output1 = net.get_superfeatures(im1_tensor.to(device), scales=[scales[scale_id]])
-        feats1 = output1[0][0]
-        attns1 = output1[1][0]
-        strenghts1 = output1[2][0]
-
-        output2 = net.get_superfeatures(im2_tensor.to(device), scales=[scales[scale_id]])
-        feats2 = output2[0][0]
-        attns2 = output2[1][0]
-        strenghts2 = output2[2][0]
-
+        # output1 = net.get_superfeatures(im1_tensor.to(device), scales=[scales[scale_id]])
+        # feats1 = output1[0][0]
+        # attns1 = output1[1][0]
+        # strenghts1 = output1[2][0]
+
+        # output2 = net.get_superfeatures(im2_tensor.to(device), scales=[scales[scale_id]])
+        # feats2 = output2[0][0]
+        # attns2 = output2[1][0]
+        # strenghts2 = output2[2][0]
+    outputs = []
+    for im_tensor in loader:
+        outputs.append(net.get_superfeatures(im_tensor.to(device), scales=[scales[scale_id]]))
+    feats1 = outputs[0][0][0]
+    attns1 = outputs[0][1][0]
+    strenghts1 = outputs[0][2][0]
+    feats2 = outputs[1][0][0]
+    attns2 = outputs[1][1][0]
+    strenghts2 = outputs[1][2][0]
     print(feats1.shape, feats2.shape)
     print(attns1.shape, attns2.shape)
     print(strenghts1.shape, strenghts2.shape)
@@ -195,7 +229,7 @@ iface = gr.Interface(
     examples=[
         ["chateau_1.png", "chateau_2.png", 2, 100, '55,14,5,4,52,57,40,9'],
         ["anafi1.jpeg", "anafi2.jpeg", 4, 50, '99,100,142,213,236'],
-        ["areopoli1.jpeg", "areopoli2.jpeg", 4, 50, '99,100,142,213,236'],
+        ["areopoli1.jpeg", "areopoli2.jpeg", 4, 50, '72,44,142,213,236'],
     ]
 )
 iface.launch(enable_queue=True)