vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/imatch-toolbox/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))
+        }

vismatch/third_party/imatch-toolbox/third_party/d2net/lib/exceptions.py

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

vismatch/third_party/imatch-toolbox/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

vismatch/third_party/imatch-toolbox/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
+        }