update: d2net lib

third_party/d2net/lib/dataset.py

@@ -0,0 +1,239 @@
+import h5py
+
+import numpy as np
+
+from PIL import Image
+
+import os
+
+import torch
+from torch.utils.data import Dataset
+
+import time
+
+from tqdm import tqdm
+
+from lib.utils import preprocess_image
+
+
+class MegaDepthDataset(Dataset):
+    def __init__(
+            self,
+            scene_list_path='megadepth_utils/train_scenes.txt',
+            scene_info_path='/local/dataset/megadepth/scene_info',
+            base_path='/local/dataset/megadepth',
+            train=True,
+            preprocessing=None,
+            min_overlap_ratio=.5,
+            max_overlap_ratio=1,
+            max_scale_ratio=np.inf,
+            pairs_per_scene=100,
+            image_size=256
+    ):
+        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):
+        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.scene_info_path, '%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:
+                continue
+            
+            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)
+        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):
+        (
+            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)
+
+        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/d2net/lib/exceptions.py

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

third_party/d2net/lib/loss.py

@@ -0,0 +1,340 @@
+import matplotlib
+import matplotlib.pyplot as plt
+
+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
+):
+    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
+
+        if plot and batch['batch_idx'] % batch['log_interval'] == 0:
+            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('train_vis/%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/d2net/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/d2net/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=True):
+        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/d2net/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/d2net/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)