add: rord libs

.gitignore

@@ -1,8 +1,6 @@
 build/
-
-lib/
+# lib
 bin/
-
 cmake_modules/
 cmake-build-debug/
 .idea/

third_party/RoRD/lib/dataloaders/datasetPhotoTourism_combined.py

@@ -0,0 +1,77 @@
+
+import os
+import time
+import random
+
+import h5py
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+import joblib
+
+import torch
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+
+from lib.utils import preprocess_image
+from lib.utils import preprocess_image, grid_positions, upscale_positions
+from lib.dataloaders.datasetPhotoTourism_ipr import PhotoTourismIPR
+from lib.dataloaders.datasetPhotoTourism_real import PhotoTourism
+
+from sys import exit, argv
+import cv2
+import csv
+
+np.random.seed(0)
+
+
+class PhotoTourismCombined(Dataset):
+    def __init__(self, base_path, preprocessing, ipr_pref=0.5, train=True, cropSize=256):
+        self.base_path = base_path
+        self.preprocessing = preprocessing
+        self.cropSize=cropSize
+
+        self.ipr_pref = ipr_pref
+
+        # self.dataset_len = 0
+        # self.dataset_len2 = 0
+
+        print("[INFO] Building Original Dataset")
+        self.PTReal = PhotoTourism(base_path, preprocessing=preprocessing, train=train, image_size=cropSize)
+        self.PTReal.build_dataset()
+
+        # self.dataset_len1 = len(self.PTReal)
+        # print("size 1:",len(self.PTReal))
+        # for _ in self.PTReal:
+        #     pass
+        # print("size 2:",len(self.PTReal))
+        self.dataset_len1 = len(self.PTReal)
+        # joblib.dump(self.PTReal.dataset, os.path.join(self.base_path, "orig_PT_2.gz"), 3)
+
+        print("[INFO] Building IPR Dataset")
+        self.PTipr = PhotoTourismIPR(base_path, preprocessing=preprocessing, train=train, cropSize=cropSize)
+        self.PTipr.build_dataset()
+
+        # self.dataset_len2 = len(self.PTipr)
+        # print("size 1:",len(self.PTipr))
+        # for _ in self.PTipr:
+        #     pass
+        # print("size 2:",len(self.PTipr))
+        self.dataset_len2 = len(self.PTipr)
+
+        # joblib.dump((self.PTipr.dataset_H, self.PTipr.valid_images), os.path.join(self.base_path, "ipr_PT_2.gz"), 3)
+
+    def __getitem__(self, idx):
+        if random.random()<self.ipr_pref:
+            return (self.PTipr[idx%self.dataset_len1], 1)
+        return (self.PTReal[idx%self.dataset_len2], 0)
+
+    def __len__(self):
+        return self.dataset_len2+self.dataset_len1
+
+
+if __name__=="__main__":
+    pt = PhotoTourismCombined("/scratch/udit/phototourism/", 'caffe', 256)
+    dl = DataLoader(pt, batch_size=1, num_workers=2)
+    for _ in dl:
+        pass

third_party/RoRD/lib/dataloaders/datasetPhotoTourism_ipr.py

@@ -0,0 +1,170 @@
+import os
+from sys import exit, argv
+import csv
+import random
+
+import joblib
+import numpy as np
+import cv2
+from PIL import Image
+from tqdm import tqdm
+
+import torch
+from torch.utils.data import Dataset
+
+from lib.utils import preprocess_image, grid_positions, upscale_positions
+
+np.random.seed(0)
+
+
+class PhotoTourismIPR(Dataset):
+	def __init__(self, base_path, preprocessing, train=True, cropSize=256):
+		self.base_path = base_path
+		self.train = train
+		self.preprocessing = preprocessing
+		self.valid_images = []
+		self.cropSize=cropSize
+
+	def getImageFiles(self):
+		img_files = []
+		img_path = "dense/images"
+		if self.train:
+			print("Inside training!!")
+
+			with open(os.path.join("configs", "train_scenes_small.txt")) as f:
+				scenes = f.read().strip("\n").split("\n")
+
+		print("[INFO]",scenes)
+		for scene in scenes:
+			image_dir = os.path.join(self.base_path, scene, img_path)
+			img_names = os.listdir(image_dir)
+			img_files += [os.path.join(image_dir, img) for img in img_names]
+		return img_files
+
+	def imgCrop(self, img1):
+		w, h = img1.size
+		left = np.random.randint(low = 0, high = w - (self.cropSize))
+		upper = np.random.randint(low = 0, high = h - (self.cropSize))
+
+		cropImg = img1.crop((left, upper, left+self.cropSize, upper+self.cropSize))
+		
+		return cropImg
+
+	def getGrid(self, im1, im2, H, scaling_steps=3):
+		h1, w1 = int(im1.shape[0]/(2**scaling_steps)), int(im1.shape[1]/(2**scaling_steps))
+		device = torch.device("cpu")
+
+		fmap_pos1 = grid_positions(h1, w1, device)
+		pos1 = upscale_positions(fmap_pos1, scaling_steps=scaling_steps).data.cpu().numpy()
+
+		pos1[[0, 1]] = pos1[[1, 0]]
+		
+		ones = np.ones((1, pos1.shape[1]))
+		pos1Homo = np.vstack((pos1, ones))
+		pos2Homo = np.dot(H, pos1Homo)
+		pos2Homo = pos2Homo/pos2Homo[2, :]
+		pos2 = pos2Homo[0:2, :]
+
+		pos1[[0, 1]] = pos1[[1, 0]]
+		pos2[[0, 1]] = pos2[[1, 0]]
+		pos1 = pos1.astype(np.float32)
+		pos2 = pos2.astype(np.float32)
+
+		ids = []
+		for i in range(pos2.shape[1]):
+			x, y = pos2[:, i]
+
+			if(2 < x < (im1.shape[0]-2) and 2 < y < (im1.shape[1]-2)):
+				ids.append(i)
+		pos1 = pos1[:, ids]
+		pos2 = pos2[:, ids]
+
+		return pos1, pos2
+	
+	def imgRotH(self, img1, min=0, max=360):
+		width, height = img1.size
+		theta = np.random.randint(low=min, high=max) * (np.pi / 180)
+		Tx = width / 2
+		Ty = height / 2
+		sx = random.uniform(-1e-2, 1e-2)
+		sy = random.uniform(-1e-2, 1e-2)
+		p1 = random.uniform(-1e-4, 1e-4)
+		p2 = random.uniform(-1e-4, 1e-4)
+
+		alpha = np.cos(theta)
+		beta = np.sin(theta)
+
+		He = np.matrix([[alpha, beta, Tx * (1 - alpha) - Ty * beta], [-beta, alpha, beta * Tx + (1 - alpha) * Ty], [0, 0, 1]])
+		Ha = np.matrix([[1, sy, 0], [sx, 1, 0], [0, 0, 1]])
+		Hp = np.matrix([[1, 0, 0], [0, 1, 0], [p1, p2, 1]])
+
+		H = He @ Ha @ Hp
+
+		return H, theta
+
+	def build_dataset(self):
+		print("Building Dataset.")
+
+		imgFiles = self.getImageFiles()
+
+		for idx in tqdm(range(len(imgFiles))):
+
+			img = imgFiles[idx]
+			img1 = Image.open(img)
+
+			if(img1.mode != 'RGB'):
+				img1 = img1.convert('RGB')
+			if(img1.size[0] < self.cropSize or img1.size[1] < self.cropSize):
+				continue
+
+			self.valid_images.append(img)
+
+	def __len__(self):
+		return len(self.valid_images)
+
+	def __getitem__(self, idx):
+		while 1:
+			try:
+				img = self.valid_images[idx]    
+				   
+				img1 = Image.open(img)
+				img1 = self.imgCrop(img1)
+				width, height = img1.size
+
+				H, theta = self.imgRotH(img1, min=0, max=360)
+
+				img1 = np.array(img1)
+				img2 = cv2.warpPerspective(img1, H, dsize=(width,height))
+				img2 = np.array(img2)
+
+				pos1, pos2 =  self.getGrid(img1, img2, H)
+
+				assert (len(pos1) != 0 and len(pos2) != 0)
+				break
+			except IndexError:
+				print("IndexError")
+				exit(1)
+			except:
+				del self.valid_images[idx]
+
+		img1 = preprocess_image(img1, preprocessing=self.preprocessing)
+		img2 = preprocess_image(img2, preprocessing=self.preprocessing)
+
+		return {
+			'image1': torch.from_numpy(img1.astype(np.float32)),
+			'image2': torch.from_numpy(img2.astype(np.float32)),
+			'pos1': torch.from_numpy(pos1.astype(np.float32)),
+			'pos2': torch.from_numpy(pos2.astype(np.float32)),
+			'H': np.array(H),
+			'theta': np.array([theta])
+		}
+
+
+if __name__ == '__main__':
+	rootDir = argv[1]
+
+	training_dataset = PhotoTourismIPR(rootDir, 'caffe')
+	training_dataset.build_dataset()
+
+	data = training_dataset[0]
+	print(data['image1'].shape, data['image2'].shape, data['pos1'].shape, data['pos2'].shape, len(training_dataset))

third_party/RoRD/lib/dataloaders/datasetPhotoTourism_real.py

@@ -0,0 +1,258 @@
+
+import os
+import time
+from tqdm import tqdm
+
+import h5py
+import numpy as np
+from PIL import Image
+
+import torch
+from torch.utils.data import Dataset
+from lib.utils import preprocess_image
+
+import joblib
+
+
+class PhotoTourism(Dataset):
+    def __init__(
+            self,
+            #scene_list_path='megadepth_utils/train_scenes.txt',
+            # scene_info_path='/local/dataset/megadepth/scene_info',
+            base_path='/scratch/udit/phototourism',
+            train=True,
+            preprocessing=None,
+            min_overlap_ratio=.5,
+            max_overlap_ratio=1,
+            max_scale_ratio=np.inf,
+            pairs_per_scene=500,
+            image_size=256
+    ):
+        if train:
+            scene_list_path = os.path.join(base_path, "train_scenes.txt.bkp")
+        else:
+            scene_list_path = os.path.join(base_path, "valid_scenes.txt")
+        self.scenes = []
+        with open(scene_list_path, 'r') as f:
+            lines = f.readlines()
+            for line in lines:
+                self.scenes.append(line.strip('\n'))
+
+        # self.scene_info_path = scene_info_path
+        self.base_path = base_path
+
+        self.train = train
+
+        self.preprocessing = preprocessing
+
+        self.min_overlap_ratio = min_overlap_ratio
+        self.max_overlap_ratio = max_overlap_ratio
+        self.max_scale_ratio = max_scale_ratio
+
+        self.pairs_per_scene = pairs_per_scene
+
+        self.image_size = image_size
+
+        self.dataset = []
+
+    def build_dataset(self):
+        cache_path = os.path.join(self.base_path, "orig_PT_2.gz")
+        if os.path.exists(cache_path):
+            self.dataset = joblib.load(cache_path)
+            return
+
+        self.dataset = []
+        if not self.train:
+            np_random_state = np.random.get_state()
+            np.random.seed(42)
+            print('Building the validation dataset...')
+        else:
+            print('Building a new training dataset...')
+        
+        for scene in tqdm(self.scenes, total=len(self.scenes)):
+            
+            scene_info_path = os.path.join(
+                self.base_path, scene, '%s.npz' % scene
+            )
+            
+            if not os.path.exists(scene_info_path):
+                continue
+            
+            scene_info = np.load(scene_info_path, allow_pickle=True)
+            overlap_matrix = scene_info['overlap_matrix']
+            scale_ratio_matrix = scene_info['scale_ratio_matrix']
+
+            valid =  np.logical_and(
+                np.logical_and(
+                    overlap_matrix >= self.min_overlap_ratio,
+                    overlap_matrix <= self.max_overlap_ratio
+                ),
+                scale_ratio_matrix <= self.max_scale_ratio
+            )
+
+            pairs = np.vstack(np.where(valid))
+            try:
+                selected_ids = np.random.choice(
+                    pairs.shape[1], self.pairs_per_scene
+                )
+            except:
+                return
+
+            image_paths = scene_info['image_paths']
+            depth_paths = scene_info['depth_paths']
+            points3D_id_to_2D = scene_info['points3D_id_to_2D']
+            points3D_id_to_ndepth = scene_info['points3D_id_to_ndepth']
+            intrinsics = scene_info['intrinsics']
+            poses = scene_info['poses']
+
+            for pair_idx in selected_ids:
+                idx1 = pairs[0, pair_idx]
+                idx2 = pairs[1, pair_idx]
+                matches = np.array(list(
+                    points3D_id_to_2D[idx1].keys() &
+                    points3D_id_to_2D[idx2].keys()
+                ))
+
+                # Scale filtering
+                matches_nd1 = np.array([points3D_id_to_ndepth[idx1][match] for match in matches])
+                matches_nd2 = np.array([points3D_id_to_ndepth[idx2][match] for match in matches])
+                scale_ratio = np.maximum(matches_nd1 / matches_nd2, matches_nd2 / matches_nd1)
+                matches = matches[np.where(scale_ratio <= self.max_scale_ratio)[0]]
+
+                point3D_id = np.random.choice(matches)
+                point2D1 = points3D_id_to_2D[idx1][point3D_id]
+                point2D2 = points3D_id_to_2D[idx2][point3D_id]
+                nd1 = points3D_id_to_ndepth[idx1][point3D_id]
+                nd2 = points3D_id_to_ndepth[idx2][point3D_id]
+                central_match = np.array([
+                    point2D1[1], point2D1[0],
+                    point2D2[1], point2D2[0]
+                ])
+                self.dataset.append({
+                    'image_path1': image_paths[idx1],
+                    'depth_path1': depth_paths[idx1],
+                    'intrinsics1': intrinsics[idx1],
+                    'pose1': poses[idx1],
+                    'image_path2': image_paths[idx2],
+                    'depth_path2': depth_paths[idx2],
+                    'intrinsics2': intrinsics[idx2],
+                    'pose2': poses[idx2],
+                    'central_match': central_match,
+                    'scale_ratio': max(nd1 / nd2, nd2 / nd1)
+                })
+        np.random.shuffle(self.dataset)
+        joblib.dump(self.dataset, cache_path, 3)
+        if not self.train:
+            np.random.set_state(np_random_state)
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def recover_pair(self, pair_metadata):
+        depth_path1 = os.path.join(
+            self.base_path, pair_metadata['depth_path1']
+        )
+        with h5py.File(depth_path1, 'r') as hdf5_file:
+            depth1 = np.array(hdf5_file['/depth'])
+        assert(np.min(depth1) >= 0)
+        image_path1 = os.path.join(
+            self.base_path, pair_metadata['image_path1']
+        )
+        image1 = Image.open(image_path1)
+        if image1.mode != 'RGB':
+            image1 = image1.convert('RGB')
+        image1 = np.array(image1)
+        assert(image1.shape[0] == depth1.shape[0] and image1.shape[1] == depth1.shape[1])
+        intrinsics1 = pair_metadata['intrinsics1']
+        pose1 = pair_metadata['pose1']
+
+        depth_path2 = os.path.join(
+            self.base_path, pair_metadata['depth_path2']
+        )
+        with h5py.File(depth_path2, 'r') as hdf5_file:
+            depth2 = np.array(hdf5_file['/depth'])
+        assert(np.min(depth2) >= 0)
+        image_path2 = os.path.join(
+            self.base_path, pair_metadata['image_path2']
+        )
+        image2 = Image.open(image_path2)
+        if image2.mode != 'RGB':
+            image2 = image2.convert('RGB')
+        image2 = np.array(image2)
+        assert(image2.shape[0] == depth2.shape[0] and image2.shape[1] == depth2.shape[1])
+        intrinsics2 = pair_metadata['intrinsics2']
+        pose2 = pair_metadata['pose2']
+
+        central_match = pair_metadata['central_match']
+        image1, bbox1, image2, bbox2 = self.crop(image1, image2, central_match)
+
+        depth1 = depth1[
+            bbox1[0] : bbox1[0] + self.image_size,
+            bbox1[1] : bbox1[1] + self.image_size
+        ]
+        depth2 = depth2[
+            bbox2[0] : bbox2[0] + self.image_size,
+            bbox2[1] : bbox2[1] + self.image_size
+        ]
+
+        return (
+            image1, depth1, intrinsics1, pose1, bbox1,
+            image2, depth2, intrinsics2, pose2, bbox2
+        )
+
+    def crop(self, image1, image2, central_match):
+        bbox1_i = max(int(central_match[0]) - self.image_size // 2, 0)
+        if bbox1_i + self.image_size >= image1.shape[0]:
+            bbox1_i = image1.shape[0] - self.image_size
+        bbox1_j = max(int(central_match[1]) - self.image_size // 2, 0)
+        if bbox1_j + self.image_size >= image1.shape[1]:
+            bbox1_j = image1.shape[1] - self.image_size
+
+        bbox2_i = max(int(central_match[2]) - self.image_size // 2, 0)
+        if bbox2_i + self.image_size >= image2.shape[0]:
+            bbox2_i = image2.shape[0] - self.image_size
+        bbox2_j = max(int(central_match[3]) - self.image_size // 2, 0)
+        if bbox2_j + self.image_size >= image2.shape[1]:
+            bbox2_j = image2.shape[1] - self.image_size
+
+        return (
+            image1[
+                bbox1_i : bbox1_i + self.image_size,
+                bbox1_j : bbox1_j + self.image_size
+            ],
+            np.array([bbox1_i, bbox1_j]),
+            image2[
+                bbox2_i : bbox2_i + self.image_size,
+                bbox2_j : bbox2_j + self.image_size
+            ],
+            np.array([bbox2_i, bbox2_j])
+        )
+
+    def __getitem__(self, idx):
+        while 1:
+            try:
+                (
+                    image1, depth1, intrinsics1, pose1, bbox1,
+                    image2, depth2, intrinsics2, pose2, bbox2
+                ) = self.recover_pair(self.dataset[idx])
+                image1 = preprocess_image(image1, preprocessing=self.preprocessing)
+                image2 = preprocess_image(image2, preprocessing=self.preprocessing)
+                assert np.all(image1.shape==image2.shape)
+                break
+            except IndexError:
+                idx-=1
+            except:
+                del self.dataset[idx]
+
+        return {
+            'image1': torch.from_numpy(image1.astype(np.float32)),
+            'depth1': torch.from_numpy(depth1.astype(np.float32)),
+            'intrinsics1': torch.from_numpy(intrinsics1.astype(np.float32)),
+            'pose1': torch.from_numpy(pose1.astype(np.float32)),
+            'bbox1': torch.from_numpy(bbox1.astype(np.float32)),
+            'image2': torch.from_numpy(image2.astype(np.float32)),
+            'depth2': torch.from_numpy(depth2.astype(np.float32)),
+            'intrinsics2': torch.from_numpy(intrinsics2.astype(np.float32)),
+            'pose2': torch.from_numpy(pose2.astype(np.float32)),
+            'bbox2': torch.from_numpy(bbox2.astype(np.float32))
+        }

third_party/RoRD/lib/exceptions.py

@@ -0,0 +1,6 @@
+class EmptyTensorError(Exception):
+    pass
+
+
+class NoGradientError(Exception):
+    pass

third_party/RoRD/lib/extractMatchTop.py

@@ -0,0 +1,361 @@
+import argparse
+import numpy as np
+import imageio
+import torch
+from tqdm import tqdm
+import time
+import scipy
+import scipy.io
+import scipy.misc
+
+from lib.model_test import D2Net
+from lib.utils import preprocess_image
+from lib.pyramid import process_multiscale
+
+import cv2
+import matplotlib.pyplot as plt
+import os
+from sys import exit, argv
+from PIL import Image
+from skimage.feature import match_descriptors
+from skimage.measure import ransac
+from skimage.transform import ProjectiveTransform, AffineTransform
+import pydegensac
+
+
+def extractSingle(image, model, device):
+
+	with torch.no_grad():
+		keypoints, scores, descriptors = process_multiscale(
+			image.to(device).unsqueeze(0),
+			model,
+			scales=[1]
+		)
+
+	keypoints = keypoints[:, [1, 0, 2]]
+
+	feat = {}
+	feat['keypoints'] = keypoints
+	feat['scores'] = scores
+	feat['descriptors'] = descriptors
+
+	return feat
+
+
+def siftMatching(img1, img2, HFile1, HFile2, device):
+	if HFile1 is not None:
+		H1 = np.load(HFile1)
+		H2 = np.load(HFile2)
+
+	rgbFile1 = img1
+	img1 = Image.open(img1)
+	
+	if(img1.mode != 'RGB'):
+		img1 = img1.convert('RGB')
+	img1 = np.array(img1)
+
+	if HFile1 is not None:
+		img1 = cv2.warpPerspective(img1, H1, dsize=(400,400))
+
+	#### Visualization ####
+	# cv2.imshow("Image", cv2.cvtColor(img1, cv2.COLOR_BGR2RGB))
+	# cv2.waitKey(0)
+
+	rgbFile2 = img2
+	img2 = Image.open(img2)
+	
+	if(img2.mode != 'RGB'):
+		img2 = img2.convert('RGB')
+	img2 = np.array(img2)
+
+	if HFile2 is not None:
+		img2 = cv2.warpPerspective(img2, H2, dsize=(400,400))
+
+	#### Visualization ####
+	# cv2.imshow("Image", cv2.cvtColor(img2, cv2.COLOR_BGR2RGB))
+	# cv2.waitKey(0)
+
+	# surf = cv2.xfeatures2d.SURF_create(100) # SURF
+	surf = cv2.xfeatures2d.SIFT_create()
+
+	kp1, des1 = surf.detectAndCompute(img1, None)
+	kp2, des2 = surf.detectAndCompute(img2, None)
+
+	matches = mnn_matcher(
+			torch.from_numpy(des1).float().to(device=device),
+			torch.from_numpy(des2).float().to(device=device)
+		)
+
+	src_pts = np.float32([ kp1[m[0]].pt for m in matches ]).reshape(-1, 2)
+	dst_pts = np.float32([ kp2[m[1]].pt for m in matches ]).reshape(-1, 2)
+
+	if(src_pts.shape[0] < 5 or dst_pts.shape[0] < 5):
+		return [], []
+
+	H, inliers = pydegensac.findHomography(src_pts, dst_pts, 8.0, 0.99, 10000)
+
+	n_inliers = np.sum(inliers)
+
+	inlier_keypoints_left = [cv2.KeyPoint(point[0], point[1], 1) for point in src_pts[inliers]]
+	inlier_keypoints_right = [cv2.KeyPoint(point[0], point[1], 1) for point in dst_pts[inliers]]
+	placeholder_matches = [cv2.DMatch(idx, idx, 1) for idx in range(n_inliers)]
+
+	#### Visualization ####
+	image3 = cv2.drawMatches(img1, inlier_keypoints_left, img2, inlier_keypoints_right, placeholder_matches, None)
+	image3 = cv2.cvtColor(image3, cv2.COLOR_BGR2RGB)
+	# cv2.imshow('Matches', image3)
+	# cv2.waitKey()
+
+	src_pts = np.float32([ inlier_keypoints_left[m.queryIdx].pt for m in placeholder_matches ]).reshape(-1, 2)
+	dst_pts = np.float32([ inlier_keypoints_right[m.trainIdx].pt for m in placeholder_matches ]).reshape(-1, 2)
+	
+	if HFile1 is None:
+		return src_pts, dst_pts, image3, image3
+	
+	orgSrc, orgDst = orgKeypoints(src_pts, dst_pts, H1, H2)
+	matchImg = drawOrg(cv2.imread(rgbFile1), cv2.imread(rgbFile2), orgSrc, orgDst)
+
+	return orgSrc, orgDst, matchImg, image3
+
+
+def orgKeypoints(src_pts, dst_pts, H1, H2):
+	ones = np.ones((src_pts.shape[0], 1))
+
+	src_pts = np.hstack((src_pts, ones))
+	dst_pts = np.hstack((dst_pts, ones))
+
+	orgSrc = np.linalg.inv(H1) @ src_pts.T
+	orgDst = np.linalg.inv(H2) @ dst_pts.T
+
+	orgSrc = orgSrc/orgSrc[2, :]
+	orgDst = orgDst/orgDst[2, :]
+
+	orgSrc = np.asarray(orgSrc)[0:2, :]
+	orgDst = np.asarray(orgDst)[0:2, :]
+
+	return orgSrc, orgDst
+
+
+def drawOrg(image1, image2, orgSrc, orgDst):
+	img1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)
+	img2 = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB)
+
+	for i in range(orgSrc.shape[1]):
+		im1 = cv2.circle(img1, (int(orgSrc[0, i]), int(orgSrc[1, i])), 3, (0, 0, 255), 1)
+	for i in range(orgDst.shape[1]):
+		im2 = cv2.circle(img2, (int(orgDst[0, i]), int(orgDst[1, i])), 3, (0, 0, 255), 1)
+
+	im4 = cv2.hconcat([im1, im2])
+	for i in range(orgSrc.shape[1]):
+		im4 = cv2.line(im4, (int(orgSrc[0, i]), int(orgSrc[1, i])), (int(orgDst[0, i]) +  im1.shape[1], int(orgDst[1, i])), (0, 255, 0), 1)
+	im4 = cv2.cvtColor(im4, cv2.COLOR_BGR2RGB)
+	# cv2.imshow("Image", im4)
+	# cv2.waitKey(0)
+
+	return im4
+
+
+
+def getPerspKeypoints(rgbFile1, rgbFile2, HFile1, HFile2, model, device):
+	if HFile1 is None:
+		igp1, img1 = read_and_process_image(rgbFile1, H=None)
+	else:
+		H1 = np.load(HFile1)
+		igp1, img1 = read_and_process_image(rgbFile1, H=H1)
+
+	c,h,w = igp1.shape
+
+	if HFile2 is None:
+		igp2, img2 = read_and_process_image(rgbFile2, H=None)
+	else:
+		H2 = np.load(HFile2)
+		igp2, img2 = read_and_process_image(rgbFile2, H=H2)
+
+	feat1 = extractSingle(igp1, model, device)
+	feat2 = extractSingle(igp2, model, device)
+
+	matches = mnn_matcher(
+			torch.from_numpy(feat1['descriptors']).to(device=device),
+			torch.from_numpy(feat2['descriptors']).to(device=device),
+		)
+	pos_a = feat1["keypoints"][matches[:, 0], : 2]
+	pos_b = feat2["keypoints"][matches[:, 1], : 2]
+
+	H, inliers = pydegensac.findHomography(pos_a, pos_b, 8.0, 0.99, 10000)
+	pos_a = pos_a[inliers]
+	pos_b = pos_b[inliers]
+
+	inlier_keypoints_left = [cv2.KeyPoint(point[0], point[1], 1) for point in pos_a]
+	inlier_keypoints_right = [cv2.KeyPoint(point[0], point[1], 1) for point in pos_b]
+	placeholder_matches = [cv2.DMatch(idx, idx, 1) for idx in range(len(pos_a))]
+
+	image3 = cv2.drawMatches(img1, inlier_keypoints_left, img2, inlier_keypoints_right, placeholder_matches, None, matchColor=[0, 255, 0])
+	image3 = cv2.cvtColor(image3, cv2.COLOR_BGR2RGB)
+
+	#### Visualization ####
+	# cv2.imshow('Matches', image3)
+	# cv2.waitKey()
+
+	if HFile1 is None:
+		return pos_a, pos_b, image3, image3
+
+	orgSrc, orgDst = orgKeypoints(pos_a, pos_b, H1, H2)
+	matchImg = drawOrg(cv2.imread(rgbFile1), cv2.imread(rgbFile2), orgSrc, orgDst) # Reproject matches to perspective View
+
+	return orgSrc, orgDst, matchImg, image3
+
+	
+
+###### Ensemble
+def read_and_process_image(img_path, resize=None, H=None, h=None, w=None, preprocessing='caffe'):
+	img1 = Image.open(img_path)
+	if resize:
+		img1 = img1.resize(resize)
+	if(img1.mode != 'RGB'):
+		img1 = img1.convert('RGB')
+	img1 = np.array(img1)
+	if H is not None:
+		img1 = cv2.warpPerspective(img1, H, dsize=(400, 400))
+		# cv2.imshow("Image", cv2.cvtColor(img1, cv2.COLOR_BGR2RGB))
+		# cv2.waitKey(0)
+	igp1 = torch.from_numpy(preprocess_image(img1, preprocessing=preprocessing).astype(np.float32))
+	return igp1, img1
+
+def mnn_matcher_scorer(descriptors_a, descriptors_b, k=np.inf):
+	device = descriptors_a.device
+	sim = descriptors_a @ descriptors_b.t()
+	val1, nn12 = torch.max(sim, dim=1)
+	val2, nn21 = torch.max(sim, dim=0)
+	ids1 = torch.arange(0, sim.shape[0], device=device)
+	mask = (ids1 == nn21[nn12])
+	matches = torch.stack([ids1[mask], nn12[mask]]).t()
+	remaining_matches_dist = val1[mask]
+	return matches, remaining_matches_dist
+
+def mnn_matcher(descriptors_a, descriptors_b):
+	device = descriptors_a.device
+	sim = descriptors_a @ descriptors_b.t()
+	nn12 = torch.max(sim, dim=1)[1]
+	nn21 = torch.max(sim, dim=0)[1]
+	ids1 = torch.arange(0, sim.shape[0], device=device)
+	mask = (ids1 == nn21[nn12])
+	matches = torch.stack([ids1[mask], nn12[mask]])
+	return matches.t().data.cpu().numpy()
+
+
+def getPerspKeypointsEnsemble(model1, model2, rgbFile1, rgbFile2, HFile1, HFile2, device):
+	if HFile1 is None:
+		igp1, img1 = read_and_process_image(rgbFile1, H=None)
+	else:
+		H1 = np.load(HFile1)
+		igp1, img1 = read_and_process_image(rgbFile1, H=H1)
+
+	c,h,w = igp1.shape
+
+	if HFile2 is None:
+		igp2, img2 = read_and_process_image(rgbFile2, H=None)
+	else:
+		H2 = np.load(HFile2)
+		igp2, img2 = read_and_process_image(rgbFile2, H=H2)
+
+	with torch.no_grad():
+		keypoints_a1, scores_a1, descriptors_a1 = process_multiscale(
+			igp1.to(device).unsqueeze(0),
+			model1,
+			scales=[1]
+		)
+		keypoints_a1 = keypoints_a1[:, [1, 0, 2]]
+
+		keypoints_a2, scores_a2, descriptors_a2 = process_multiscale(
+			igp1.to(device).unsqueeze(0),
+			model2,
+			scales=[1]
+		)
+		keypoints_a2 = keypoints_a2[:, [1, 0, 2]]
+
+		keypoints_b1, scores_b1, descriptors_b1 = process_multiscale(
+			igp2.to(device).unsqueeze(0),
+			model1,
+			scales=[1]
+		)
+		keypoints_b1 = keypoints_b1[:, [1, 0, 2]]
+
+		keypoints_b2, scores_b2, descriptors_b2 = process_multiscale(
+			igp2.to(device).unsqueeze(0),
+			model2,
+			scales=[1]
+		)
+		keypoints_b2 = keypoints_b2[:, [1, 0, 2]]
+
+	# calculating matches for both models
+	matches1, dist_1 = mnn_matcher_scorer(
+		torch.from_numpy(descriptors_a1).to(device=device),
+		torch.from_numpy(descriptors_b1).to(device=device),
+#                 len(matches1)
+	)
+	matches2, dist_2 = mnn_matcher_scorer(
+		torch.from_numpy(descriptors_a2).to(device=device),
+		torch.from_numpy(descriptors_b2).to(device=device),
+#                 len(matches1)
+	)
+
+	full_matches = torch.cat([matches1, matches2])
+	full_dist = torch.cat([dist_1, dist_2])
+	assert len(full_dist)==(len(dist_1)+len(dist_2)), "something wrong"
+
+	k_final = len(full_dist)//2
+	# k_final = len(full_dist)
+	# k_final = max(len(dist_1), len(dist_2))
+	top_k_mask = torch.topk(full_dist, k=k_final)[1]
+	first = []
+	second = []
+
+	for valid_id in top_k_mask:
+		if valid_id<len(dist_1):
+			first.append(valid_id)
+		else:
+			second.append(valid_id-len(dist_1))
+	# final_matches = full_matches[top_k_mask]
+
+	matches1 = matches1[torch.tensor(first, device=device).long()].data.cpu().numpy()
+	matches2 = matches2[torch.tensor(second, device=device).long()].data.cpu().numpy()
+
+	pos_a1 = keypoints_a1[matches1[:, 0], : 2]
+	pos_b1 = keypoints_b1[matches1[:, 1], : 2]
+
+	pos_a2 = keypoints_a2[matches2[:, 0], : 2]
+	pos_b2 = keypoints_b2[matches2[:, 1], : 2]
+
+	pos_a = np.concatenate([pos_a1, pos_a2], 0)
+	pos_b = np.concatenate([pos_b1, pos_b2], 0)
+
+	# pos_a, pos_b, inliers = apply_ransac(pos_a, pos_b)
+	H, inliers = pydegensac.findHomography(pos_a, pos_b, 8.0, 0.99, 10000)
+	pos_a = pos_a[inliers]
+	pos_b = pos_b[inliers]
+
+	inlier_keypoints_left = [cv2.KeyPoint(point[0], point[1], 1) for point in pos_a]
+	inlier_keypoints_right = [cv2.KeyPoint(point[0], point[1], 1) for point in pos_b]
+	placeholder_matches = [cv2.DMatch(idx, idx, 1) for idx in range(len(pos_a))]
+
+	image3 = cv2.drawMatches(img1, inlier_keypoints_left, img2, inlier_keypoints_right, placeholder_matches, None, matchColor=[0, 255, 0])
+	image3 = cv2.cvtColor(image3, cv2.COLOR_BGR2RGB)
+	# cv2.imshow('Matches', image3)
+	# cv2.waitKey()
+
+
+	orgSrc, orgDst = orgKeypoints(pos_a, pos_b, H1, H2)
+	matchImg = drawOrg(cv2.imread(rgbFile1), cv2.imread(rgbFile2), orgSrc, orgDst)
+
+	return orgSrc, orgDst, matchImg, image3
+
+
+if __name__ == '__main__':
+	WEIGHTS = '../models/rord.pth'
+	
+	srcR = argv[1]
+	trgR = argv[2]
+	srcH = argv[3]
+	trgH = argv[4]
+
+	orgSrc, orgDst = getPerspKeypoints(srcR, trgR, srcH, trgH, WEIGHTS, ('gpu'))

third_party/RoRD/lib/loss.py

@@ -0,0 +1,342 @@
+import matplotlib
+import matplotlib.pyplot as plt
+import os
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+from lib.utils import (
+    grid_positions,
+    upscale_positions,
+    downscale_positions,
+    savefig,
+    imshow_image
+)
+from lib.exceptions import NoGradientError, EmptyTensorError
+
+matplotlib.use('Agg')
+
+
+def loss_function(
+        model, batch, device, margin=1, safe_radius=4, scaling_steps=3, plot=False, plot_path=None
+):
+    output = model({
+        'image1': batch['image1'].to(device),
+        'image2': batch['image2'].to(device)
+    })
+
+    loss = torch.tensor(np.array([0], dtype=np.float32), device=device)
+    has_grad = False
+
+    n_valid_samples = 0
+    for idx_in_batch in range(batch['image1'].size(0)):
+        # Annotations
+        depth1 = batch['depth1'][idx_in_batch].to(device)  # [h1, w1]
+        intrinsics1 = batch['intrinsics1'][idx_in_batch].to(device)  # [3, 3]
+        pose1 = batch['pose1'][idx_in_batch].view(4, 4).to(device)  # [4, 4]
+        bbox1 = batch['bbox1'][idx_in_batch].to(device)  # [2]
+
+        depth2 = batch['depth2'][idx_in_batch].to(device)
+        intrinsics2 = batch['intrinsics2'][idx_in_batch].to(device)
+        pose2 = batch['pose2'][idx_in_batch].view(4, 4).to(device)
+        bbox2 = batch['bbox2'][idx_in_batch].to(device)
+
+        # Network output
+        dense_features1 = output['dense_features1'][idx_in_batch]
+        c, h1, w1 = dense_features1.size()
+        scores1 = output['scores1'][idx_in_batch].view(-1)
+
+        dense_features2 = output['dense_features2'][idx_in_batch]
+        _, h2, w2 = dense_features2.size()
+        scores2 = output['scores2'][idx_in_batch]
+
+        all_descriptors1 = F.normalize(dense_features1.view(c, -1), dim=0)
+        descriptors1 = all_descriptors1
+
+        all_descriptors2 = F.normalize(dense_features2.view(c, -1), dim=0)
+
+        # Warp the positions from image 1 to image 2
+        fmap_pos1 = grid_positions(h1, w1, device)
+        pos1 = upscale_positions(fmap_pos1, scaling_steps=scaling_steps)
+        try:
+            pos1, pos2, ids = warp(
+                pos1,
+                depth1, intrinsics1, pose1, bbox1,
+                depth2, intrinsics2, pose2, bbox2
+            )
+        except EmptyTensorError:
+            continue
+        fmap_pos1 = fmap_pos1[:, ids]
+        descriptors1 = descriptors1[:, ids]
+        scores1 = scores1[ids]
+
+        # Skip the pair if not enough GT correspondences are available
+        if ids.size(0) < 128:
+            continue
+
+        # Descriptors at the corresponding positions
+        fmap_pos2 = torch.round(
+            downscale_positions(pos2, scaling_steps=scaling_steps)
+        ).long()
+        descriptors2 = F.normalize(
+            dense_features2[:, fmap_pos2[0, :], fmap_pos2[1, :]],
+            dim=0
+        )
+        positive_distance = 2 - 2 * (
+            descriptors1.t().unsqueeze(1) @ descriptors2.t().unsqueeze(2)
+        ).squeeze()
+
+        all_fmap_pos2 = grid_positions(h2, w2, device)
+        position_distance = torch.max(
+            torch.abs(
+                fmap_pos2.unsqueeze(2).float() -
+                all_fmap_pos2.unsqueeze(1)
+            ),
+            dim=0
+        )[0]
+        is_out_of_safe_radius = position_distance > safe_radius
+        distance_matrix = 2 - 2 * (descriptors1.t() @ all_descriptors2)
+        negative_distance2 = torch.min(
+            distance_matrix + (1 - is_out_of_safe_radius.float()) * 10.,
+            dim=1
+        )[0]
+
+        all_fmap_pos1 = grid_positions(h1, w1, device)
+        position_distance = torch.max(
+            torch.abs(
+                fmap_pos1.unsqueeze(2).float() -
+                all_fmap_pos1.unsqueeze(1)
+            ),
+            dim=0
+        )[0]
+        is_out_of_safe_radius = position_distance > safe_radius
+        distance_matrix = 2 - 2 * (descriptors2.t() @ all_descriptors1)
+        negative_distance1 = torch.min(
+            distance_matrix + (1 - is_out_of_safe_radius.float()) * 10.,
+            dim=1
+        )[0]
+
+        diff = positive_distance - torch.min(
+            negative_distance1, negative_distance2
+        )
+
+        scores2 = scores2[fmap_pos2[0, :], fmap_pos2[1, :]]
+
+        loss = loss + (
+            torch.sum(scores1 * scores2 * F.relu(margin + diff)) /
+            torch.sum(scores1 * scores2)
+        )
+
+        has_grad = True
+        n_valid_samples += 1
+
+        # print(plot, batch['batch_idx'],batch['log_interval'])
+        if plot and batch['batch_idx'] % batch['log_interval'] == 0:
+            # print("should plot")
+            pos1_aux = pos1.cpu().numpy()
+            pos2_aux = pos2.cpu().numpy()
+            k = pos1_aux.shape[1]
+            col = np.random.rand(k, 3)
+            n_sp = 4
+            plt.figure()
+            plt.subplot(1, n_sp, 1)
+            im1 = imshow_image(
+                batch['image1'][idx_in_batch].cpu().numpy(),
+                preprocessing=batch['preprocessing']
+            )
+            plt.imshow(im1)
+            plt.scatter(
+                pos1_aux[1, :], pos1_aux[0, :],
+                s=0.25**2, c=col, marker=',', alpha=0.5
+            )
+            plt.axis('off')
+            plt.subplot(1, n_sp, 2)
+            plt.imshow(
+                output['scores1'][idx_in_batch].data.cpu().numpy(),
+                cmap='Reds'
+            )
+            plt.axis('off')
+            plt.subplot(1, n_sp, 3)
+            im2 = imshow_image(
+                batch['image2'][idx_in_batch].cpu().numpy(),
+                preprocessing=batch['preprocessing']
+            )
+            plt.imshow(im2)
+            plt.scatter(
+                pos2_aux[1, :], pos2_aux[0, :],
+                s=0.25**2, c=col, marker=',', alpha=0.5
+            )
+            plt.axis('off')
+            plt.subplot(1, n_sp, 4)
+            plt.imshow(
+                output['scores2'][idx_in_batch].data.cpu().numpy(),
+                cmap='Reds'
+            )
+            plt.axis('off')
+            savefig(os.path.join(plot_path, '%s.%02d.%02d.%d.png' % (
+                'train' if batch['train'] else 'valid',
+                batch['epoch_idx'],
+                batch['batch_idx'] // batch['log_interval'],
+                idx_in_batch
+            )), dpi=300)
+            plt.close()
+
+    if not has_grad:
+        raise NoGradientError
+
+    loss = loss / n_valid_samples
+
+    return loss
+
+
+def interpolate_depth(pos, depth):
+    device = pos.device
+
+    ids = torch.arange(0, pos.size(1), device=device)
+
+    h, w = depth.size()
+
+    i = pos[0, :]
+    j = pos[1, :]
+
+    # Valid corners
+    i_top_left = torch.floor(i).long()
+    j_top_left = torch.floor(j).long()
+    valid_top_left = torch.min(i_top_left >= 0, j_top_left >= 0)
+
+    i_top_right = torch.floor(i).long()
+    j_top_right = torch.ceil(j).long()
+    valid_top_right = torch.min(i_top_right >= 0, j_top_right < w)
+
+    i_bottom_left = torch.ceil(i).long()
+    j_bottom_left = torch.floor(j).long()
+    valid_bottom_left = torch.min(i_bottom_left < h, j_bottom_left >= 0)
+
+    i_bottom_right = torch.ceil(i).long()
+    j_bottom_right = torch.ceil(j).long()
+    valid_bottom_right = torch.min(i_bottom_right < h, j_bottom_right < w)
+
+    valid_corners = torch.min(
+        torch.min(valid_top_left, valid_top_right),
+        torch.min(valid_bottom_left, valid_bottom_right)
+    )
+
+    i_top_left = i_top_left[valid_corners]
+    j_top_left = j_top_left[valid_corners]
+
+    i_top_right = i_top_right[valid_corners]
+    j_top_right = j_top_right[valid_corners]
+
+    i_bottom_left = i_bottom_left[valid_corners]
+    j_bottom_left = j_bottom_left[valid_corners]
+
+    i_bottom_right = i_bottom_right[valid_corners]
+    j_bottom_right = j_bottom_right[valid_corners]
+
+    ids = ids[valid_corners]
+    if ids.size(0) == 0:
+        raise EmptyTensorError
+
+    # Valid depth
+    valid_depth = torch.min(
+        torch.min(
+            depth[i_top_left, j_top_left] > 0,
+            depth[i_top_right, j_top_right] > 0
+        ),
+        torch.min(
+            depth[i_bottom_left, j_bottom_left] > 0,
+            depth[i_bottom_right, j_bottom_right] > 0
+        )
+    )
+
+    i_top_left = i_top_left[valid_depth]
+    j_top_left = j_top_left[valid_depth]
+
+    i_top_right = i_top_right[valid_depth]
+    j_top_right = j_top_right[valid_depth]
+
+    i_bottom_left = i_bottom_left[valid_depth]
+    j_bottom_left = j_bottom_left[valid_depth]
+
+    i_bottom_right = i_bottom_right[valid_depth]
+    j_bottom_right = j_bottom_right[valid_depth]
+
+    ids = ids[valid_depth]
+    if ids.size(0) == 0:
+        raise EmptyTensorError
+
+    # Interpolation
+    i = i[ids]
+    j = j[ids]
+    dist_i_top_left = i - i_top_left.float()
+    dist_j_top_left = j - j_top_left.float()
+    w_top_left = (1 - dist_i_top_left) * (1 - dist_j_top_left)
+    w_top_right = (1 - dist_i_top_left) * dist_j_top_left
+    w_bottom_left = dist_i_top_left * (1 - dist_j_top_left)
+    w_bottom_right = dist_i_top_left * dist_j_top_left
+
+    interpolated_depth = (
+        w_top_left * depth[i_top_left, j_top_left] +
+        w_top_right * depth[i_top_right, j_top_right] +
+        w_bottom_left * depth[i_bottom_left, j_bottom_left] +
+        w_bottom_right * depth[i_bottom_right, j_bottom_right]
+    )
+
+    pos = torch.cat([i.view(1, -1), j.view(1, -1)], dim=0)
+
+    return [interpolated_depth, pos, ids]
+
+
+def uv_to_pos(uv):
+    return torch.cat([uv[1, :].view(1, -1), uv[0, :].view(1, -1)], dim=0)
+
+
+def warp(
+        pos1,
+        depth1, intrinsics1, pose1, bbox1,
+        depth2, intrinsics2, pose2, bbox2
+):
+    device = pos1.device
+
+    Z1, pos1, ids = interpolate_depth(pos1, depth1)
+
+    # COLMAP convention
+    u1 = pos1[1, :] + bbox1[1] + .5
+    v1 = pos1[0, :] + bbox1[0] + .5
+
+    X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
+    Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])
+
+    XYZ1_hom = torch.cat([
+        X1.view(1, -1),
+        Y1.view(1, -1),
+        Z1.view(1, -1),
+        torch.ones(1, Z1.size(0), device=device)
+    ], dim=0)
+    XYZ2_hom = torch.chain_matmul(pose2, torch.inverse(pose1), XYZ1_hom)
+    XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)
+
+    uv2_hom = torch.matmul(intrinsics2, XYZ2)
+    uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)
+
+    u2 = uv2[0, :] - bbox2[1] - .5
+    v2 = uv2[1, :] - bbox2[0] - .5
+    uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)
+
+    annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)
+
+    ids = ids[new_ids]
+    pos1 = pos1[:, new_ids]
+    estimated_depth = XYZ2[2, new_ids]
+
+    inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05
+
+    ids = ids[inlier_mask]
+    if ids.size(0) == 0:
+        raise EmptyTensorError
+
+    pos2 = pos2[:, inlier_mask]
+    pos1 = pos1[:, inlier_mask]
+
+    return pos1, pos2, ids

third_party/RoRD/lib/losses/lossPhotoTourism.py

@@ -0,0 +1,232 @@
+import matplotlib
+import matplotlib.pyplot as plt
+
+import numpy as np
+import cv2
+from sys import exit
+
+import torch
+import torch.nn.functional as F
+
+from lib.utils import (
+	grid_positions,
+	upscale_positions,
+	downscale_positions,
+	savefig,
+	imshow_image
+)
+from lib.exceptions import NoGradientError, EmptyTensorError
+
+matplotlib.use('Agg')
+
+
+def loss_function(
+		model, batch, device, margin=1, safe_radius=4, scaling_steps=3, plot=False, plot_path=None
+):
+	output = model({
+		'image1': batch['image1'].to(device),
+		'image2': batch['image2'].to(device)
+	})
+	
+	
+	loss = torch.tensor(np.array([0], dtype=np.float32), device=device)
+	has_grad = False
+
+	n_valid_samples = 0
+	for idx_in_batch in range(batch['image1'].size(0)):
+		# Network output
+		dense_features1 = output['dense_features1'][idx_in_batch]
+		c, h1, w1 = dense_features1.size()
+		scores1 = output['scores1'][idx_in_batch].view(-1)
+
+		dense_features2 = output['dense_features2'][idx_in_batch]
+		_, h2, w2 = dense_features2.size()
+		scores2 = output['scores2'][idx_in_batch]
+
+		all_descriptors1 = F.normalize(dense_features1.view(c, -1), dim=0)
+		descriptors1 = all_descriptors1
+
+		all_descriptors2 = F.normalize(dense_features2.view(c, -1), dim=0)
+
+		fmap_pos1 = grid_positions(h1, w1, device)
+
+		pos1 = batch['pos1'][idx_in_batch].to(device)
+		pos2 = batch['pos2'][idx_in_batch].to(device)
+
+		ids = idsAlign(pos1, device, h1, w1)
+
+		fmap_pos1 = fmap_pos1[:, ids]
+		descriptors1 = descriptors1[:, ids]
+		scores1 = scores1[ids]
+
+		# Skip the pair if not enough GT correspondences are available
+		if ids.size(0) < 128:
+			continue
+
+		# Descriptors at the corresponding positions
+		fmap_pos2 = torch.round(
+			downscale_positions(pos2, scaling_steps=scaling_steps)
+		).long()
+
+		descriptors2 = F.normalize(
+			dense_features2[:, fmap_pos2[0, :], fmap_pos2[1, :]],
+			dim=0
+		)
+		positive_distance = 2 - 2 * (
+			descriptors1.t().unsqueeze(1) @ descriptors2.t().unsqueeze(2)
+		).squeeze()
+
+		all_fmap_pos2 = grid_positions(h2, w2, device)
+		position_distance = torch.max(
+			torch.abs(
+				fmap_pos2.unsqueeze(2).float() -
+				all_fmap_pos2.unsqueeze(1)
+			),
+			dim=0
+		)[0]
+		is_out_of_safe_radius = position_distance > safe_radius
+
+		distance_matrix = 2 - 2 * (descriptors1.t() @ all_descriptors2)
+
+		negative_distance2 = torch.min(
+			distance_matrix + (1 - is_out_of_safe_radius.float()) * 10.,
+			dim=1
+		)[0]
+
+		all_fmap_pos1 = grid_positions(h1, w1, device)
+		position_distance = torch.max(
+			torch.abs(
+				fmap_pos1.unsqueeze(2).float() -
+				all_fmap_pos1.unsqueeze(1)
+			),
+			dim=0
+		)[0]
+		is_out_of_safe_radius = position_distance > safe_radius
+
+		distance_matrix = 2 - 2 * (descriptors2.t() @ all_descriptors1)
+
+		negative_distance1 = torch.min(
+			distance_matrix + (1 - is_out_of_safe_radius.float()) * 10.,
+			dim=1
+		)[0]
+
+		diff = positive_distance - torch.min(
+			negative_distance1, negative_distance2
+		)
+
+		scores2 = scores2[fmap_pos2[0, :], fmap_pos2[1, :]]
+
+		loss = loss + (
+			torch.sum(scores1 * scores2 * F.relu(margin + diff)) /
+			(torch.sum(scores1 * scores2) )
+		)
+
+		has_grad = True
+		n_valid_samples += 1
+
+		if plot and batch['batch_idx'] % batch['log_interval'] == 0:
+			drawTraining(batch['image1'], batch['image2'], pos1, pos2, batch, idx_in_batch, output, save=True, plot_path=plot_path)
+
+	if not has_grad:
+		raise NoGradientError
+
+	loss = loss / (n_valid_samples )
+
+	return loss
+
+
+def idsAlign(pos1, device, h1, w1):
+	pos1D = downscale_positions(pos1, scaling_steps=3)
+	row = pos1D[0, :]
+	col = pos1D[1, :]
+
+	ids = []
+
+	for i in range(row.shape[0]):
+
+		index = ((w1) * (row[i])) + (col[i])
+		ids.append(index)
+
+	ids = torch.round(torch.Tensor(ids)).long().to(device)
+
+	return ids
+
+
+def drawTraining(image1, image2, pos1, pos2, batch, idx_in_batch, output, save=False, plot_path="train_viz"):
+	pos1_aux = pos1.cpu().numpy()
+	pos2_aux = pos2.cpu().numpy()
+
+	k = pos1_aux.shape[1]
+	col = np.random.rand(k, 3)
+	n_sp = 4
+	plt.figure()
+	plt.subplot(1, n_sp, 1)
+	im1 = imshow_image(
+		image1[0].cpu().numpy(),
+		preprocessing=batch['preprocessing']
+	)
+	plt.imshow(im1)
+	plt.scatter(
+		pos1_aux[1, :], pos1_aux[0, :],
+		s=0.25**2, c=col, marker=',', alpha=0.5
+	)
+	plt.axis('off')
+	plt.subplot(1, n_sp, 2)
+	plt.imshow(
+		output['scores1'][idx_in_batch].data.cpu().numpy(),
+		cmap='Reds'
+	)
+	plt.axis('off')
+	plt.subplot(1, n_sp, 3)
+	im2 = imshow_image(
+		image2[0].cpu().numpy(),
+		preprocessing=batch['preprocessing']
+	)
+	plt.imshow(im2)
+	plt.scatter(
+		pos2_aux[1, :], pos2_aux[0, :],
+		s=0.25**2, c=col, marker=',', alpha=0.5
+	)
+	plt.axis('off')
+	plt.subplot(1, n_sp, 4)
+	plt.imshow(
+		output['scores2'][idx_in_batch].data.cpu().numpy(),
+		cmap='Reds'
+	)
+	plt.axis('off')
+
+	if(save == True):
+		savefig(plot_path+'/%s.%02d.%02d.%d.png' % (
+			'train' if batch['train'] else 'valid',
+			batch['epoch_idx'],
+			batch['batch_idx'] // batch['log_interval'],
+			idx_in_batch
+		), dpi=300)
+	else:
+		plt.show()
+
+	plt.close()
+
+	im1 = cv2.cvtColor(im1, cv2.COLOR_BGR2RGB)
+	im2 = cv2.cvtColor(im2, cv2.COLOR_BGR2RGB)
+
+	for i in range(0, pos1_aux.shape[1], 5):
+		im1 = cv2.circle(im1, (pos1_aux[1, i], pos1_aux[0, i]), 1, (0, 0, 255), 2)
+	for i in range(0, pos2_aux.shape[1], 5):
+		im2 = cv2.circle(im2, (pos2_aux[1, i], pos2_aux[0, i]), 1, (0, 0, 255), 2)
+
+	im3 = cv2.hconcat([im1, im2])
+
+	for i in range(0, pos1_aux.shape[1], 5):
+		im3 = cv2.line(im3, (int(pos1_aux[1, i]), int(pos1_aux[0, i])), (int(pos2_aux[1, i]) +  im1.shape[1], int(pos2_aux[0, i])), (0, 255, 0), 1)
+
+	if(save == True):
+		cv2.imwrite(plot_path+'/%s.%02d.%02d.%d.png' % (
+			'train_corr' if batch['train'] else 'valid',
+			batch['epoch_idx'],
+			batch['batch_idx'] // batch['log_interval'],
+			idx_in_batch
+		), im3)
+	else:
+		cv2.imshow('Image', im3)
+		cv2.waitKey(0)

third_party/RoRD/lib/model.py

@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torchvision.models as models
+
+
+class DenseFeatureExtractionModule(nn.Module):
+    def __init__(self, finetune_feature_extraction=False, use_cuda=True):
+        super(DenseFeatureExtractionModule, self).__init__()
+
+        model = models.vgg16()
+        vgg16_layers = [
+            'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2',
+            'pool1',
+            'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2',
+            'pool2',
+            'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3',
+            'pool3',
+            'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3',
+            'pool4',
+            'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3',
+            'pool5'
+        ]
+        conv4_3_idx = vgg16_layers.index('conv4_3')
+
+        self.model = nn.Sequential(
+            *list(model.features.children())[: conv4_3_idx + 1]
+        )
+
+        self.num_channels = 512
+
+        # Fix forward parameters
+        for param in self.model.parameters():
+            param.requires_grad = False
+        if finetune_feature_extraction:
+            # Unlock conv4_3
+            for param in list(self.model.parameters())[-2 :]:
+                param.requires_grad = True
+
+        if use_cuda:
+            self.model = self.model.cuda()
+
+    def forward(self, batch):
+        output = self.model(batch)
+        return output
+
+
+class SoftDetectionModule(nn.Module):
+    def __init__(self, soft_local_max_size=3):
+        super(SoftDetectionModule, self).__init__()
+
+        self.soft_local_max_size = soft_local_max_size
+
+        self.pad = self.soft_local_max_size // 2
+
+    def forward(self, batch):
+        b = batch.size(0)
+
+        batch = F.relu(batch)
+
+        max_per_sample = torch.max(batch.view(b, -1), dim=1)[0]
+        exp = torch.exp(batch / max_per_sample.view(b, 1, 1, 1))
+        sum_exp = (
+            self.soft_local_max_size ** 2 *
+            F.avg_pool2d(
+                F.pad(exp, [self.pad] * 4, mode='constant', value=1.),
+                self.soft_local_max_size, stride=1
+            )
+        )
+        local_max_score = exp / sum_exp
+
+        depth_wise_max = torch.max(batch, dim=1)[0]
+        depth_wise_max_score = batch / depth_wise_max.unsqueeze(1)
+
+        all_scores = local_max_score * depth_wise_max_score
+        score = torch.max(all_scores, dim=1)[0]
+
+        score = score / torch.sum(score.view(b, -1), dim=1).view(b, 1, 1)
+
+        return score
+
+
+class D2Net(nn.Module):
+    def __init__(self, model_file=None, use_cuda=True):
+        super(D2Net, self).__init__()
+
+        self.dense_feature_extraction = DenseFeatureExtractionModule(
+            finetune_feature_extraction=True,
+            use_cuda=use_cuda
+        )
+
+        self.detection = SoftDetectionModule()
+
+        if model_file is not None:
+            if use_cuda:
+                self.load_state_dict(torch.load(model_file)['model'])
+            else:
+                self.load_state_dict(torch.load(model_file, map_location='cpu')['model'])
+
+    def forward(self, batch):
+        b = batch['image1'].size(0)
+
+        dense_features = self.dense_feature_extraction(
+            torch.cat([batch['image1'], batch['image2']], dim=0)
+        )
+
+        scores = self.detection(dense_features)
+
+        dense_features1 = dense_features[: b, :, :, :]
+        dense_features2 = dense_features[b :, :, :, :]
+
+        scores1 = scores[: b, :, :]
+        scores2 = scores[b :, :, :]
+
+        return {
+            'dense_features1': dense_features1,
+            'scores1': scores1,
+            'dense_features2': dense_features2,
+            'scores2': scores2
+        }

third_party/RoRD/lib/model_test.py

@@ -0,0 +1,187 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DenseFeatureExtractionModule(nn.Module):
+    def __init__(self, use_relu=True, use_cuda=True):
+        super(DenseFeatureExtractionModule, self).__init__()
+
+        self.model = nn.Sequential(
+            nn.Conv2d(3, 64, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, stride=2),
+            nn.Conv2d(64, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, stride=2),
+            nn.Conv2d(128, 256, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.AvgPool2d(2, stride=1),
+            nn.Conv2d(256, 512, 3, padding=2, dilation=2),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, 512, 3, padding=2, dilation=2),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, 512, 3, padding=2, dilation=2),
+        )
+        self.num_channels = 512
+
+        self.use_relu = use_relu
+
+        if use_cuda:
+            self.model = self.model.cuda()
+
+    def forward(self, batch):
+        output = self.model(batch)
+        if self.use_relu:
+            output = F.relu(output)
+        return output
+
+
+class D2Net(nn.Module):
+    def __init__(self, model_file=None, use_relu=True, use_cuda=False):
+        super(D2Net, self).__init__()
+
+        self.dense_feature_extraction = DenseFeatureExtractionModule(
+            use_relu=use_relu, use_cuda=use_cuda
+        )
+
+        self.detection = HardDetectionModule()
+
+        self.localization = HandcraftedLocalizationModule()
+
+        if model_file is not None:
+            if use_cuda:
+                self.load_state_dict(torch.load(model_file)['model'])
+            else:
+                self.load_state_dict(torch.load(model_file, map_location='cpu')['model'])
+
+    def forward(self, batch):
+        _, _, h, w = batch.size()
+        dense_features = self.dense_feature_extraction(batch)
+
+        detections = self.detection(dense_features)
+
+        displacements = self.localization(dense_features)
+
+        return {
+            'dense_features': dense_features,
+            'detections': detections,
+            'displacements': displacements
+        }
+
+
+class HardDetectionModule(nn.Module):
+    def __init__(self, edge_threshold=5):
+        super(HardDetectionModule, self).__init__()
+
+        self.edge_threshold = edge_threshold
+
+        self.dii_filter = torch.tensor(
+            [[0, 1., 0], [0, -2., 0], [0, 1., 0]]
+        ).view(1, 1, 3, 3)
+        self.dij_filter = 0.25 * torch.tensor(
+            [[1., 0, -1.], [0, 0., 0], [-1., 0, 1.]]
+        ).view(1, 1, 3, 3)
+        self.djj_filter = torch.tensor(
+            [[0, 0, 0], [1., -2., 1.], [0, 0, 0]]
+        ).view(1, 1, 3, 3)
+
+    def forward(self, batch):
+        b, c, h, w = batch.size()
+        device = batch.device
+
+        depth_wise_max = torch.max(batch, dim=1)[0]
+        is_depth_wise_max = (batch == depth_wise_max)
+        del depth_wise_max
+
+        local_max = F.max_pool2d(batch, 3, stride=1, padding=1)
+        is_local_max = (batch == local_max)
+        del local_max
+
+        dii = F.conv2d(
+            batch.view(-1, 1, h, w), self.dii_filter.to(device), padding=1
+        ).view(b, c, h, w)
+        dij = F.conv2d(
+            batch.view(-1, 1, h, w), self.dij_filter.to(device), padding=1
+        ).view(b, c, h, w)
+        djj = F.conv2d(
+            batch.view(-1, 1, h, w), self.djj_filter.to(device), padding=1
+        ).view(b, c, h, w)
+
+        det = dii * djj - dij * dij
+        tr = dii + djj
+        del dii, dij, djj
+
+        threshold = (self.edge_threshold + 1) ** 2 / self.edge_threshold
+        is_not_edge = torch.min(tr * tr / det <= threshold, det > 0)
+
+        detected = torch.min(
+            is_depth_wise_max,
+            torch.min(is_local_max, is_not_edge)
+        )
+        del is_depth_wise_max, is_local_max, is_not_edge
+
+        return detected
+
+
+class HandcraftedLocalizationModule(nn.Module):
+    def __init__(self):
+        super(HandcraftedLocalizationModule, self).__init__()
+
+        self.di_filter = torch.tensor(
+            [[0, -0.5, 0], [0, 0, 0], [0,  0.5, 0]]
+        ).view(1, 1, 3, 3)
+        self.dj_filter = torch.tensor(
+            [[0, 0, 0], [-0.5, 0, 0.5], [0, 0, 0]]
+        ).view(1, 1, 3, 3)
+
+        self.dii_filter = torch.tensor(
+            [[0, 1., 0], [0, -2., 0], [0, 1., 0]]
+        ).view(1, 1, 3, 3)
+        self.dij_filter = 0.25 * torch.tensor(
+            [[1., 0, -1.], [0, 0., 0], [-1., 0, 1.]]
+        ).view(1, 1, 3, 3)
+        self.djj_filter = torch.tensor(
+            [[0, 0, 0], [1., -2., 1.], [0, 0, 0]]
+        ).view(1, 1, 3, 3)
+
+    def forward(self, batch):
+        b, c, h, w = batch.size()
+        device = batch.device
+
+        dii = F.conv2d(
+            batch.view(-1, 1, h, w), self.dii_filter.to(device), padding=1
+        ).view(b, c, h, w)
+        dij = F.conv2d(
+            batch.view(-1, 1, h, w), self.dij_filter.to(device), padding=1
+        ).view(b, c, h, w)
+        djj = F.conv2d(
+            batch.view(-1, 1, h, w), self.djj_filter.to(device), padding=1
+        ).view(b, c, h, w)
+        det = dii * djj - dij * dij
+
+        inv_hess_00 = djj / det
+        inv_hess_01 = -dij / det
+        inv_hess_11 = dii / det
+        del dii, dij, djj, det
+
+        di = F.conv2d(
+            batch.view(-1, 1, h, w), self.di_filter.to(device), padding=1
+        ).view(b, c, h, w)
+        dj = F.conv2d(
+            batch.view(-1, 1, h, w), self.dj_filter.to(device), padding=1
+        ).view(b, c, h, w)
+
+        step_i = -(inv_hess_00 * di + inv_hess_01 * dj)
+        step_j = -(inv_hess_01 * di + inv_hess_11 * dj)
+        del inv_hess_00, inv_hess_01, inv_hess_11, di, dj
+
+        return torch.stack([step_i, step_j], dim=1)

third_party/RoRD/lib/pyramid.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from lib.exceptions import EmptyTensorError
+from lib.utils import interpolate_dense_features, upscale_positions
+
+
+def process_multiscale(image, model, scales=[.5, 1, 2]):
+    b, _, h_init, w_init = image.size()
+    device = image.device
+    assert(b == 1)
+
+    all_keypoints = torch.zeros([3, 0])
+    all_descriptors = torch.zeros([
+        model.dense_feature_extraction.num_channels, 0
+    ])
+    all_scores = torch.zeros(0)
+
+    previous_dense_features = None
+    banned = None
+    for idx, scale in enumerate(scales):
+        current_image = F.interpolate(
+            image, scale_factor=scale,
+            mode='bilinear', align_corners=True
+        )
+        _, _, h_level, w_level = current_image.size()
+
+        dense_features = model.dense_feature_extraction(current_image)
+        del current_image
+
+        _, _, h, w = dense_features.size()
+
+        # Sum the feature maps.
+        if previous_dense_features is not None:
+            dense_features += F.interpolate(
+                previous_dense_features, size=[h, w],
+                mode='bilinear', align_corners=True
+            )
+            del previous_dense_features
+
+        # Recover detections.
+        detections = model.detection(dense_features)
+        if banned is not None:
+            banned = F.interpolate(banned.float(), size=[h, w]).bool()
+            detections = torch.min(detections, ~banned)
+            banned = torch.max(
+                torch.max(detections, dim=1)[0].unsqueeze(1), banned
+            )
+        else:
+            banned = torch.max(detections, dim=1)[0].unsqueeze(1)
+        fmap_pos = torch.nonzero(detections[0].cpu()).t()
+        del detections
+
+        # Recover displacements.
+        displacements = model.localization(dense_features)[0].cpu()
+        displacements_i = displacements[
+            0, fmap_pos[0, :], fmap_pos[1, :], fmap_pos[2, :]
+        ]
+        displacements_j = displacements[
+            1, fmap_pos[0, :], fmap_pos[1, :], fmap_pos[2, :]
+        ]
+        del displacements
+
+        mask = torch.min(
+            torch.abs(displacements_i) < 0.5,
+            torch.abs(displacements_j) < 0.5
+        )
+        fmap_pos = fmap_pos[:, mask]
+        valid_displacements = torch.stack([
+            displacements_i[mask],
+            displacements_j[mask]
+        ], dim=0)
+        del mask, displacements_i, displacements_j
+
+        fmap_keypoints = fmap_pos[1 :, :].float() + valid_displacements
+        del valid_displacements
+
+        try:
+            raw_descriptors, _, ids = interpolate_dense_features(
+                fmap_keypoints.to(device),
+                dense_features[0]
+            )
+        except EmptyTensorError:
+            continue
+        fmap_pos = fmap_pos.to(device)
+        fmap_keypoints = fmap_keypoints.to(device)
+        fmap_pos = fmap_pos[:, ids]
+        fmap_keypoints = fmap_keypoints[:, ids]
+        del ids
+
+        keypoints = upscale_positions(fmap_keypoints, scaling_steps=2)
+        del fmap_keypoints
+
+        descriptors = F.normalize(raw_descriptors, dim=0).cpu()
+        del raw_descriptors
+
+        keypoints[0, :] *= h_init / h_level
+        keypoints[1, :] *= w_init / w_level
+
+        fmap_pos = fmap_pos.cpu()
+        keypoints = keypoints.cpu()
+
+        keypoints = torch.cat([
+            keypoints,
+            torch.ones([1, keypoints.size(1)]) * 1 / scale,
+        ], dim=0)
+
+        scores = dense_features[
+            0, fmap_pos[0, :], fmap_pos[1, :], fmap_pos[2, :]
+        ].cpu() / (idx + 1)
+        del fmap_pos
+
+        all_keypoints = torch.cat([all_keypoints, keypoints], dim=1)
+        all_descriptors = torch.cat([all_descriptors, descriptors], dim=1)
+        all_scores = torch.cat([all_scores, scores], dim=0)
+        del keypoints, descriptors
+
+        previous_dense_features = dense_features
+        del dense_features
+    del previous_dense_features, banned
+
+    keypoints = all_keypoints.t().detach().numpy()
+    del all_keypoints
+    scores = all_scores.detach().numpy()
+    del all_scores
+    descriptors = all_descriptors.t().detach().numpy()
+    del all_descriptors
+    return keypoints, scores, descriptors

third_party/RoRD/lib/utils.py

@@ -0,0 +1,167 @@
+import matplotlib.pyplot as plt
+
+import numpy as np
+
+import torch
+
+from lib.exceptions import EmptyTensorError
+
+
+def preprocess_image(image, preprocessing=None):
+    image = image.astype(np.float32)
+    image = np.transpose(image, [2, 0, 1])
+    if preprocessing is None:
+        pass
+    elif preprocessing == 'caffe':
+        # RGB -> BGR
+        image = image[:: -1, :, :]
+        # Zero-center by mean pixel
+        mean = np.array([103.939, 116.779, 123.68])
+        image = image - mean.reshape([3, 1, 1])
+    elif preprocessing == 'torch':
+        image /= 255.0
+        mean = np.array([0.485, 0.456, 0.406])
+        std = np.array([0.229, 0.224, 0.225])
+        image = (image - mean.reshape([3, 1, 1])) / std.reshape([3, 1, 1])
+    else:
+        raise ValueError('Unknown preprocessing parameter.')
+    return image
+
+
+def imshow_image(image, preprocessing=None):
+    if preprocessing is None:
+        pass
+    elif preprocessing == 'caffe':
+        mean = np.array([103.939, 116.779, 123.68])
+        image = image + mean.reshape([3, 1, 1])
+        # RGB -> BGR
+        image = image[:: -1, :, :]
+    elif preprocessing == 'torch':
+        mean = np.array([0.485, 0.456, 0.406])
+        std = np.array([0.229, 0.224, 0.225])
+        image = image * std.reshape([3, 1, 1]) + mean.reshape([3, 1, 1])
+        image *= 255.0
+    else:
+        raise ValueError('Unknown preprocessing parameter.')
+    image = np.transpose(image, [1, 2, 0])
+    image = np.round(image).astype(np.uint8)
+    return image
+
+
+def grid_positions(h, w, device, matrix=False):
+    lines = torch.arange(
+        0, h, device=device
+    ).view(-1, 1).float().repeat(1, w)
+    columns = torch.arange(
+        0, w, device=device
+    ).view(1, -1).float().repeat(h, 1)
+    if matrix:
+        return torch.stack([lines, columns], dim=0)
+    else:
+        return torch.cat([lines.view(1, -1), columns.view(1, -1)], dim=0)
+
+
+def upscale_positions(pos, scaling_steps=0):
+    for _ in range(scaling_steps):
+        pos = pos * 2 + 0.5
+    return pos
+
+
+def downscale_positions(pos, scaling_steps=0):
+    for _ in range(scaling_steps):
+        pos = (pos - 0.5) / 2
+    return pos
+
+
+def interpolate_dense_features(pos, dense_features, return_corners=False):
+    device = pos.device
+
+    ids = torch.arange(0, pos.size(1), device=device)
+
+    _, h, w = dense_features.size()
+
+    i = pos[0, :]
+    j = pos[1, :]
+
+    # Valid corners
+    i_top_left = torch.floor(i).long()
+    j_top_left = torch.floor(j).long()
+    valid_top_left = torch.min(i_top_left >= 0, j_top_left >= 0)
+
+    i_top_right = torch.floor(i).long()
+    j_top_right = torch.ceil(j).long()
+    valid_top_right = torch.min(i_top_right >= 0, j_top_right < w)
+
+    i_bottom_left = torch.ceil(i).long()
+    j_bottom_left = torch.floor(j).long()
+    valid_bottom_left = torch.min(i_bottom_left < h, j_bottom_left >= 0)
+
+    i_bottom_right = torch.ceil(i).long()
+    j_bottom_right = torch.ceil(j).long()
+    valid_bottom_right = torch.min(i_bottom_right < h, j_bottom_right < w)
+
+    valid_corners = torch.min(
+        torch.min(valid_top_left, valid_top_right),
+        torch.min(valid_bottom_left, valid_bottom_right)
+    )
+
+    i_top_left = i_top_left[valid_corners]
+    j_top_left = j_top_left[valid_corners]
+
+    i_top_right = i_top_right[valid_corners]
+    j_top_right = j_top_right[valid_corners]
+
+    i_bottom_left = i_bottom_left[valid_corners]
+    j_bottom_left = j_bottom_left[valid_corners]
+
+    i_bottom_right = i_bottom_right[valid_corners]
+    j_bottom_right = j_bottom_right[valid_corners]
+
+    ids = ids[valid_corners]
+    if ids.size(0) == 0:
+        raise EmptyTensorError
+
+    # Interpolation
+    i = i[ids]
+    j = j[ids]
+    dist_i_top_left = i - i_top_left.float()
+    dist_j_top_left = j - j_top_left.float()
+    w_top_left = (1 - dist_i_top_left) * (1 - dist_j_top_left)
+    w_top_right = (1 - dist_i_top_left) * dist_j_top_left
+    w_bottom_left = dist_i_top_left * (1 - dist_j_top_left)
+    w_bottom_right = dist_i_top_left * dist_j_top_left
+
+    descriptors = (
+        w_top_left * dense_features[:, i_top_left, j_top_left] +
+        w_top_right * dense_features[:, i_top_right, j_top_right] +
+        w_bottom_left * dense_features[:, i_bottom_left, j_bottom_left] +
+        w_bottom_right * dense_features[:, i_bottom_right, j_bottom_right]
+    )
+
+    pos = torch.cat([i.view(1, -1), j.view(1, -1)], dim=0)
+
+    if not return_corners:
+        return [descriptors, pos, ids]
+    else:
+        corners = torch.stack([
+            torch.stack([i_top_left, j_top_left], dim=0),
+            torch.stack([i_top_right, j_top_right], dim=0),
+            torch.stack([i_bottom_left, j_bottom_left], dim=0),
+            torch.stack([i_bottom_right, j_bottom_right], dim=0)
+        ], dim=0)
+        return [descriptors, pos, ids, corners]
+
+
+def savefig(filepath, fig=None, dpi=None):
+    # TomNorway - https://stackoverflow.com/a/53516034
+    if not fig:
+        fig = plt.gcf()
+
+    plt.subplots_adjust(0, 0, 1, 1, 0, 0)
+    for ax in fig.axes:
+        ax.axis('off')
+        ax.margins(0, 0)
+        ax.xaxis.set_major_locator(plt.NullLocator())
+        ax.yaxis.set_major_locator(plt.NullLocator())
+
+    fig.savefig(filepath, pad_inches=0, bbox_inches='tight', dpi=dpi)