unimatch demo

app.py

@@ -0,0 +1,160 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import gradio as gr
+
+from unimatch.unimatch import UniMatch
+from utils.flow_viz import flow_to_image
+from dataloader.stereo import transforms
+from utils.visualization import vis_disparity
+
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+
+@torch.no_grad()
+def inference(image1, image2, task='flow'):
+    """Inference on an image pair for optical flow or stereo disparity prediction"""
+
+    model = UniMatch(feature_channels=128,
+                     num_scales=2,
+                     upsample_factor=4,
+                     ffn_dim_expansion=4,
+                     num_transformer_layers=6,
+                     reg_refine=True,
+                     task=task)
+
+    model.eval()
+
+    if task == 'flow':
+        checkpoint_path = 'pretrained/gmflow-scale2-regrefine6-mixdata-train320x576-4e7b215d.pth'
+    else:
+        checkpoint_path = 'pretrained/gmstereo-scale2-regrefine3-resumeflowthings-mixdata-train320x640-ft640x960-e4e291fd.pth'
+
+    checkpoint_flow = torch.load(checkpoint_path)
+    model.load_state_dict(checkpoint_flow['model'], strict=True)
+
+    padding_factor = 32
+    attn_type = 'swin' if task == 'flow' else 'self_swin2d_cross_swin1d'
+    attn_splits_list = [2, 8]
+    corr_radius_list = [-1, 4]
+    prop_radius_list = [-1, 1]
+    num_reg_refine = 6 if task == 'flow' else 3
+
+    # smaller inference size for faster speed
+    max_inference_size = [384, 768] if task == 'flow' else [640, 960]
+
+    transpose_img = False
+
+    image1 = np.array(image1).astype(np.float32)
+    image2 = np.array(image2).astype(np.float32)
+
+    if len(image1.shape) == 2:  # gray image
+        image1 = np.tile(image1[..., None], (1, 1, 3))
+        image2 = np.tile(image2[..., None], (1, 1, 3))
+    else:
+        image1 = image1[..., :3]
+        image2 = image2[..., :3]
+
+    if task == 'flow':
+        image1 = torch.from_numpy(image1).permute(2, 0, 1).float().unsqueeze(0)
+        image2 = torch.from_numpy(image2).permute(2, 0, 1).float().unsqueeze(0)
+    else:
+        val_transform_list = [transforms.ToTensor(),
+                              transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+                              ]
+
+        val_transform = transforms.Compose(val_transform_list)
+
+        sample = {'left': image1, 'right': image2}
+        sample = val_transform(sample)
+
+        image1 = sample['left'].unsqueeze(0)  # [1, 3, H, W]
+        image2 = sample['right'].unsqueeze(0)  # [1, 3, H, W]
+
+    # the model is trained with size: width > height
+    if task == 'flow' and image1.size(-2) > image1.size(-1):
+        image1 = torch.transpose(image1, -2, -1)
+        image2 = torch.transpose(image2, -2, -1)
+        transpose_img = True
+
+    nearest_size = [int(np.ceil(image1.size(-2) / padding_factor)) * padding_factor,
+                    int(np.ceil(image1.size(-1) / padding_factor)) * padding_factor]
+
+    inference_size = [min(max_inference_size[0], nearest_size[0]), min(max_inference_size[1], nearest_size[1])]
+
+    assert isinstance(inference_size, list) or isinstance(inference_size, tuple)
+    ori_size = image1.shape[-2:]
+
+    # resize before inference
+    if inference_size[0] != ori_size[0] or inference_size[1] != ori_size[1]:
+        image1 = F.interpolate(image1, size=inference_size, mode='bilinear',
+                               align_corners=True)
+        image2 = F.interpolate(image2, size=inference_size, mode='bilinear',
+                               align_corners=True)
+
+    results_dict = model(image1, image2,
+                         attn_type=attn_type,
+                         attn_splits_list=attn_splits_list,
+                         corr_radius_list=corr_radius_list,
+                         prop_radius_list=prop_radius_list,
+                         num_reg_refine=num_reg_refine,
+                         task=task,
+                         )
+
+    flow_pr = results_dict['flow_preds'][-1]  # [1, 2, H, W] or [1, H, W]
+
+    # resize back
+    if task == 'flow':
+        if inference_size[0] != ori_size[0] or inference_size[1] != ori_size[1]:
+            flow_pr = F.interpolate(flow_pr, size=ori_size, mode='bilinear',
+                                    align_corners=True)
+            flow_pr[:, 0] = flow_pr[:, 0] * ori_size[-1] / inference_size[-1]
+            flow_pr[:, 1] = flow_pr[:, 1] * ori_size[-2] / inference_size[-2]
+    else:
+        if inference_size[0] != ori_size[0] or inference_size[1] != ori_size[1]:
+            pred_disp = F.interpolate(flow_pr.unsqueeze(1), size=ori_size,
+                                      mode='bilinear',
+                                      align_corners=True).squeeze(1)  # [1, H, W]
+            pred_disp = pred_disp * ori_size[-1] / float(inference_size[-1])
+
+    if task == 'flow':
+        if transpose_img:
+            flow_pr = torch.transpose(flow_pr, -2, -1)
+
+        flow = flow_pr[0].permute(1, 2, 0).cpu().numpy()  # [H, W, 2]
+
+        output = flow_to_image(flow)  # [H, W, 3]
+    else:
+        disp = pred_disp[0].cpu().numpy()
+
+        output = vis_disparity(disp, return_rgb=True)
+
+    return Image.fromarray(output)
+
+
+title = "UniMatch"
+
+description = "<p style='text-align: center'>Optical flow and stereo matching demo for <a href='https://haofeixu.github.io/unimatch/' target='_blank'>Unifying Flow, Stereo and Depth Estimation</a> | <a href='https://arxiv.org/abs/2211.05783' target='_blank'>Paper</a> | <a href='https://github.com/autonomousvision/unimatch' target='_blank'>Code</a> | <a href='https://colab.research.google.com/drive/1r5m-xVy3Kw60U-m5VB-aQ98oqqg_6cab?usp=sharing' target='_blank'>Colab</a><br>Simply upload your images or click one of the provided examples.<br>The <strong>first three</strong> examples are video frames for <strong>flow</strong> task, and the <strong>last three</strong> are stereo pairs for <strong>stereo</strong> task.<br><strong>Select the task type according to your input images</strong>.</p>"
+
+examples = [
+    ['demo/flow_kitti_test_000197_10.png', 'demo/flow_kitti_test_000197_11.png'],
+    ['demo/flow_sintel_cave_3_frame_0049.png', 'demo/flow_sintel_cave_3_frame_0050.png'],
+    ['demo/flow_davis_skate-jump_00059.jpg', 'demo/flow_davis_skate-jump_00060.jpg'],
+    ['demo/stereo_drivingstereo_test_2018-07-11-14-48-52_2018-07-11-14-58-34-673_left.jpg',
+     'demo/stereo_drivingstereo_test_2018-07-11-14-48-52_2018-07-11-14-58-34-673_right.jpg'],
+    ['demo/stereo_middlebury_plants_im0.png', 'demo/stereo_middlebury_plants_im1.png'],
+    ['demo/stereo_holopix_left.png', 'demo/stereo_holopix_right.png']
+]
+
+gr.Interface(
+    inference,
+    [gr.Image(type="pil", label="Image1"), gr.Image(type="pil", label="Image2"), gr.Radio(choices=['flow', 'stereo'], value='flow', label='Task')],
+    gr.Image(type="pil", label="Flow/Disparity"),
+    title=title,
+    description=description,
+    examples=examples,
+    thumbnail="https://haofeixu.github.io/unimatch/resources/teaser.svg",
+    allow_flagging="auto",
+).launch(debug=True)

dataloader/stereo/transforms.py

@@ -0,0 +1,434 @@
+from __future__ import division
+import torch
+import numpy as np
+from PIL import Image
+import torchvision.transforms.functional as F
+import random
+import cv2
+
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, sample):
+        for t in self.transforms:
+            sample = t(sample)
+        return sample
+
+
+class ToTensor(object):
+    """Convert numpy array to torch tensor"""
+
+    def __init__(self, no_normalize=False):
+        self.no_normalize = no_normalize
+
+    def __call__(self, sample):
+        left = np.transpose(sample['left'], (2, 0, 1))  # [3, H, W]
+        if self.no_normalize:
+            sample['left'] = torch.from_numpy(left)
+        else:
+            sample['left'] = torch.from_numpy(left) / 255.
+        right = np.transpose(sample['right'], (2, 0, 1))
+
+        if self.no_normalize:
+            sample['right'] = torch.from_numpy(right)
+        else:
+            sample['right'] = torch.from_numpy(right) / 255.
+
+        # disp = np.expand_dims(sample['disp'], axis=0)  # [1, H, W]
+        if 'disp' in sample.keys():
+            disp = sample['disp']  # [H, W]
+            sample['disp'] = torch.from_numpy(disp)
+
+        return sample
+
+
+class Normalize(object):
+    """Normalize image, with type tensor"""
+
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, sample):
+
+        norm_keys = ['left', 'right']
+
+        for key in norm_keys:
+            # Images have converted to tensor, with shape [C, H, W]
+            for t, m, s in zip(sample[key], self.mean, self.std):
+                t.sub_(m).div_(s)
+
+        return sample
+
+
+class RandomCrop(object):
+    def __init__(self, img_height, img_width):
+        self.img_height = img_height
+        self.img_width = img_width
+
+    def __call__(self, sample):
+        ori_height, ori_width = sample['left'].shape[:2]
+
+        # pad zero when crop size is larger than original image size
+        if self.img_height > ori_height or self.img_width > ori_width:
+
+            # can be used for only pad one side
+            top_pad = max(self.img_height - ori_height, 0)
+            right_pad = max(self.img_width - ori_width, 0)
+
+            # try edge padding
+            sample['left'] = np.lib.pad(sample['left'],
+                                        ((top_pad, 0), (0, right_pad), (0, 0)),
+                                        mode='edge')
+            sample['right'] = np.lib.pad(sample['right'],
+                                         ((top_pad, 0), (0, right_pad), (0, 0)),
+                                         mode='edge')
+
+            if 'disp' in sample.keys():
+                sample['disp'] = np.lib.pad(sample['disp'],
+                                            ((top_pad, 0), (0, right_pad)),
+                                            mode='constant',
+                                            constant_values=0)
+
+            # update image resolution
+            ori_height, ori_width = sample['left'].shape[:2]
+
+        assert self.img_height <= ori_height and self.img_width <= ori_width
+
+        # Training: random crop
+        self.offset_x = np.random.randint(ori_width - self.img_width + 1)
+
+        start_height = 0
+        assert ori_height - start_height >= self.img_height
+
+        self.offset_y = np.random.randint(start_height, ori_height - self.img_height + 1)
+
+        sample['left'] = self.crop_img(sample['left'])
+        sample['right'] = self.crop_img(sample['right'])
+        if 'disp' in sample.keys():
+            sample['disp'] = self.crop_img(sample['disp'])
+
+        return sample
+
+    def crop_img(self, img):
+        return img[self.offset_y:self.offset_y + self.img_height,
+               self.offset_x:self.offset_x + self.img_width]
+
+
+class RandomVerticalFlip(object):
+    """Randomly vertically filps"""
+
+    def __call__(self, sample):
+        if np.random.random() < 0.5:
+            sample['left'] = np.copy(np.flipud(sample['left']))
+            sample['right'] = np.copy(np.flipud(sample['right']))
+
+            sample['disp'] = np.copy(np.flipud(sample['disp']))
+
+        return sample
+
+
+class ToPILImage(object):
+
+    def __call__(self, sample):
+        sample['left'] = Image.fromarray(sample['left'].astype('uint8'))
+        sample['right'] = Image.fromarray(sample['right'].astype('uint8'))
+
+        return sample
+
+
+class ToNumpyArray(object):
+
+    def __call__(self, sample):
+        sample['left'] = np.array(sample['left']).astype(np.float32)
+        sample['right'] = np.array(sample['right']).astype(np.float32)
+
+        return sample
+
+
+# Random coloring
+class RandomContrast(object):
+    """Random contrast"""
+
+    def __init__(self,
+                 asymmetric_color_aug=True,
+                 ):
+
+        self.asymmetric_color_aug = asymmetric_color_aug
+
+    def __call__(self, sample):
+        if np.random.random() < 0.5:
+            contrast_factor = np.random.uniform(0.8, 1.2)
+
+            sample['left'] = F.adjust_contrast(sample['left'], contrast_factor)
+
+            if self.asymmetric_color_aug and np.random.random() < 0.5:
+                contrast_factor = np.random.uniform(0.8, 1.2)
+
+            sample['right'] = F.adjust_contrast(sample['right'], contrast_factor)
+
+        return sample
+
+
+class RandomGamma(object):
+
+    def __init__(self,
+                 asymmetric_color_aug=True,
+                 ):
+
+        self.asymmetric_color_aug = asymmetric_color_aug
+
+    def __call__(self, sample):
+        if np.random.random() < 0.5:
+            gamma = np.random.uniform(0.7, 1.5)  # adopted from FlowNet
+
+            sample['left'] = F.adjust_gamma(sample['left'], gamma)
+
+            if self.asymmetric_color_aug and np.random.random() < 0.5:
+                gamma = np.random.uniform(0.7, 1.5)  # adopted from FlowNet
+
+            sample['right'] = F.adjust_gamma(sample['right'], gamma)
+
+        return sample
+
+
+class RandomBrightness(object):
+
+    def __init__(self,
+                 asymmetric_color_aug=True,
+                 ):
+
+        self.asymmetric_color_aug = asymmetric_color_aug
+
+    def __call__(self, sample):
+        if np.random.random() < 0.5:
+            brightness = np.random.uniform(0.5, 2.0)
+
+            sample['left'] = F.adjust_brightness(sample['left'], brightness)
+
+            if self.asymmetric_color_aug and np.random.random() < 0.5:
+                brightness = np.random.uniform(0.5, 2.0)
+
+            sample['right'] = F.adjust_brightness(sample['right'], brightness)
+
+        return sample
+
+
+class RandomHue(object):
+
+    def __init__(self,
+                 asymmetric_color_aug=True,
+                 ):
+
+        self.asymmetric_color_aug = asymmetric_color_aug
+
+    def __call__(self, sample):
+        if np.random.random() < 0.5:
+            hue = np.random.uniform(-0.1, 0.1)
+
+            sample['left'] = F.adjust_hue(sample['left'], hue)
+
+            if self.asymmetric_color_aug and np.random.random() < 0.5:
+                hue = np.random.uniform(-0.1, 0.1)
+
+            sample['right'] = F.adjust_hue(sample['right'], hue)
+
+        return sample
+
+
+class RandomSaturation(object):
+
+    def __init__(self,
+                 asymmetric_color_aug=True,
+                 ):
+
+        self.asymmetric_color_aug = asymmetric_color_aug
+
+    def __call__(self, sample):
+        if np.random.random() < 0.5:
+            saturation = np.random.uniform(0.8, 1.2)
+
+            sample['left'] = F.adjust_saturation(sample['left'], saturation)
+
+            if self.asymmetric_color_aug and np.random.random() < 0.5:
+                saturation = np.random.uniform(0.8, 1.2)
+
+            sample['right'] = F.adjust_saturation(sample['right'], saturation)
+
+        return sample
+
+
+class RandomColor(object):
+
+    def __init__(self,
+                 asymmetric_color_aug=True,
+                 ):
+
+        self.asymmetric_color_aug = asymmetric_color_aug
+
+    def __call__(self, sample):
+        transforms = [RandomContrast(asymmetric_color_aug=self.asymmetric_color_aug),
+                      RandomGamma(asymmetric_color_aug=self.asymmetric_color_aug),
+                      RandomBrightness(asymmetric_color_aug=self.asymmetric_color_aug),
+                      RandomHue(asymmetric_color_aug=self.asymmetric_color_aug),
+                      RandomSaturation(asymmetric_color_aug=self.asymmetric_color_aug)]
+
+        sample = ToPILImage()(sample)
+
+        if np.random.random() < 0.5:
+            # A single transform
+            t = random.choice(transforms)
+            sample = t(sample)
+        else:
+            # Combination of transforms
+            # Random order
+            random.shuffle(transforms)
+            for t in transforms:
+                sample = t(sample)
+
+        sample = ToNumpyArray()(sample)
+
+        return sample
+
+
+class RandomScale(object):
+    def __init__(self,
+                 min_scale=-0.4,
+                 max_scale=0.4,
+                 crop_width=512,
+                 nearest_interp=False,  # for sparse gt
+                 ):
+        self.min_scale = min_scale
+        self.max_scale = max_scale
+        self.crop_width = crop_width
+        self.nearest_interp = nearest_interp
+
+    def __call__(self, sample):
+        if np.random.rand() < 0.5:
+            h, w = sample['disp'].shape
+
+            scale_x = 2 ** np.random.uniform(self.min_scale, self.max_scale)
+
+            scale_x = np.clip(scale_x, self.crop_width / float(w), None)
+
+            # only random scale x axis
+            sample['left'] = cv2.resize(sample['left'], None, fx=scale_x, fy=1., interpolation=cv2.INTER_LINEAR)
+            sample['right'] = cv2.resize(sample['right'], None, fx=scale_x, fy=1., interpolation=cv2.INTER_LINEAR)
+
+            sample['disp'] = cv2.resize(
+                sample['disp'], None, fx=scale_x, fy=1.,
+                interpolation=cv2.INTER_LINEAR if not self.nearest_interp else cv2.INTER_NEAREST
+            ) * scale_x
+
+            if 'pseudo_disp' in sample and sample['pseudo_disp'] is not None:
+                sample['pseudo_disp'] = cv2.resize(sample['pseudo_disp'], None, fx=scale_x, fy=1.,
+                                                   interpolation=cv2.INTER_LINEAR) * scale_x
+
+        return sample
+
+
+class Resize(object):
+    def __init__(self,
+                 scale_x=1,
+                 scale_y=1,
+                 nearest_interp=True,  # for sparse gt
+                 ):
+        """
+        Resize low-resolution data to high-res for mixed dataset training
+        """
+        self.scale_x = scale_x
+        self.scale_y = scale_y
+        self.nearest_interp = nearest_interp
+
+    def __call__(self, sample):
+        scale_x = self.scale_x
+        scale_y = self.scale_y
+
+        sample['left'] = cv2.resize(sample['left'], None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
+        sample['right'] = cv2.resize(sample['right'], None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
+
+        sample['disp'] = cv2.resize(
+            sample['disp'], None, fx=scale_x, fy=scale_y,
+            interpolation=cv2.INTER_LINEAR if not self.nearest_interp else cv2.INTER_NEAREST
+        ) * scale_x
+
+        return sample
+
+
+class RandomGrayscale(object):
+    def __init__(self, p=0.2):
+        self.p = p
+
+    def __call__(self, sample):
+        if np.random.random() < self.p:
+            sample = ToPILImage()(sample)
+
+            # only supported in higher version pytorch
+            # default output channels is 1
+            sample['left'] = F.rgb_to_grayscale(sample['left'], num_output_channels=3)
+            sample['right'] = F.rgb_to_grayscale(sample['right'], num_output_channels=3)
+
+            sample = ToNumpyArray()(sample)
+
+        return sample
+
+
+class RandomRotateShiftRight(object):
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, sample):
+        if np.random.random() < self.p:
+            angle, pixel = 0.1, 2
+            px = np.random.uniform(-pixel, pixel)
+            ag = np.random.uniform(-angle, angle)
+
+            right_img = sample['right']
+
+            image_center = (
+                np.random.uniform(0, right_img.shape[0]),
+                np.random.uniform(0, right_img.shape[1])
+            )
+
+            rot_mat = cv2.getRotationMatrix2D(image_center, ag, 1.0)
+            right_img = cv2.warpAffine(
+                right_img, rot_mat, right_img.shape[1::-1], flags=cv2.INTER_LINEAR
+            )
+            trans_mat = np.float32([[1, 0, 0], [0, 1, px]])
+            right_img = cv2.warpAffine(
+                right_img, trans_mat, right_img.shape[1::-1], flags=cv2.INTER_LINEAR
+            )
+
+            sample['right'] = right_img
+
+        return sample
+
+
+class RandomOcclusion(object):
+    def __init__(self, p=0.5,
+                 occlusion_mask_zero=False):
+        self.p = p
+        self.occlusion_mask_zero = occlusion_mask_zero
+
+    def __call__(self, sample):
+        bounds = [50, 100]
+        if np.random.random() < self.p:
+            img2 = sample['right']
+            ht, wd = img2.shape[:2]
+
+            if self.occlusion_mask_zero:
+                mean_color = 0
+            else:
+                mean_color = np.mean(img2.reshape(-1, 3), axis=0)
+
+            x0 = np.random.randint(0, wd)
+            y0 = np.random.randint(0, ht)
+            dx = np.random.randint(bounds[0], bounds[1])
+            dy = np.random.randint(bounds[0], bounds[1])
+            img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color
+
+            sample['right'] = img2
+
+        return sample

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+matplotlib
+opencv-python
+pillow

unimatch/attention.py

@@ -0,0 +1,253 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import split_feature, merge_splits, split_feature_1d, merge_splits_1d
+
+
+def single_head_full_attention(q, k, v):
+    # q, k, v: [B, L, C]
+    assert q.dim() == k.dim() == v.dim() == 3
+
+    scores = torch.matmul(q, k.permute(0, 2, 1)) / (q.size(2) ** .5)  # [B, L, L]
+    attn = torch.softmax(scores, dim=2)  # [B, L, L]
+    out = torch.matmul(attn, v)  # [B, L, C]
+
+    return out
+
+
+def single_head_full_attention_1d(q, k, v,
+                                  h=None,
+                                  w=None,
+                                  ):
+    # q, k, v: [B, L, C]
+
+    assert h is not None and w is not None
+    assert q.size(1) == h * w
+
+    b, _, c = q.size()
+
+    q = q.view(b, h, w, c)  # [B, H, W, C]
+    k = k.view(b, h, w, c)
+    v = v.view(b, h, w, c)
+
+    scale_factor = c ** 0.5
+
+    scores = torch.matmul(q, k.permute(0, 1, 3, 2)) / scale_factor  # [B, H, W, W]
+
+    attn = torch.softmax(scores, dim=-1)
+
+    out = torch.matmul(attn, v).view(b, -1, c)  # [B, H*W, C]
+
+    return out
+
+
+def single_head_split_window_attention(q, k, v,
+                                       num_splits=1,
+                                       with_shift=False,
+                                       h=None,
+                                       w=None,
+                                       attn_mask=None,
+                                       ):
+    # ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
+    # q, k, v: [B, L, C]
+    assert q.dim() == k.dim() == v.dim() == 3
+
+    assert h is not None and w is not None
+    assert q.size(1) == h * w
+
+    b, _, c = q.size()
+
+    b_new = b * num_splits * num_splits
+
+    window_size_h = h // num_splits
+    window_size_w = w // num_splits
+
+    q = q.view(b, h, w, c)  # [B, H, W, C]
+    k = k.view(b, h, w, c)
+    v = v.view(b, h, w, c)
+
+    scale_factor = c ** 0.5
+
+    if with_shift:
+        assert attn_mask is not None  # compute once
+        shift_size_h = window_size_h // 2
+        shift_size_w = window_size_w // 2
+
+        q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
+        k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
+        v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
+
+    q = split_feature(q, num_splits=num_splits, channel_last=True)  # [B*K*K, H/K, W/K, C]
+    k = split_feature(k, num_splits=num_splits, channel_last=True)
+    v = split_feature(v, num_splits=num_splits, channel_last=True)
+
+    scores = torch.matmul(q.view(b_new, -1, c), k.view(b_new, -1, c).permute(0, 2, 1)
+                          ) / scale_factor  # [B*K*K, H/K*W/K, H/K*W/K]
+
+    if with_shift:
+        scores += attn_mask.repeat(b, 1, 1)
+
+    attn = torch.softmax(scores, dim=-1)
+
+    out = torch.matmul(attn, v.view(b_new, -1, c))  # [B*K*K, H/K*W/K, C]
+
+    out = merge_splits(out.view(b_new, h // num_splits, w // num_splits, c),
+                       num_splits=num_splits, channel_last=True)  # [B, H, W, C]
+
+    # shift back
+    if with_shift:
+        out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2))
+
+    out = out.view(b, -1, c)
+
+    return out
+
+
+def single_head_split_window_attention_1d(q, k, v,
+                                          relative_position_bias=None,
+                                          num_splits=1,
+                                          with_shift=False,
+                                          h=None,
+                                          w=None,
+                                          attn_mask=None,
+                                          ):
+    # q, k, v: [B, L, C]
+
+    assert h is not None and w is not None
+    assert q.size(1) == h * w
+
+    b, _, c = q.size()
+
+    b_new = b * num_splits * h
+
+    window_size_w = w // num_splits
+
+    q = q.view(b * h, w, c)  # [B*H, W, C]
+    k = k.view(b * h, w, c)
+    v = v.view(b * h, w, c)
+
+    scale_factor = c ** 0.5
+
+    if with_shift:
+        assert attn_mask is not None  # compute once
+        shift_size_w = window_size_w // 2
+
+        q = torch.roll(q, shifts=-shift_size_w, dims=1)
+        k = torch.roll(k, shifts=-shift_size_w, dims=1)
+        v = torch.roll(v, shifts=-shift_size_w, dims=1)
+
+    q = split_feature_1d(q, num_splits=num_splits)  # [B*H*K, W/K, C]
+    k = split_feature_1d(k, num_splits=num_splits)
+    v = split_feature_1d(v, num_splits=num_splits)
+
+    scores = torch.matmul(q.view(b_new, -1, c), k.view(b_new, -1, c).permute(0, 2, 1)
+                          ) / scale_factor  # [B*H*K, W/K, W/K]
+
+    if with_shift:
+        # attn_mask: [K, W/K, W/K]
+        scores += attn_mask.repeat(b * h, 1, 1)  # [B*H*K, W/K, W/K]
+
+    attn = torch.softmax(scores, dim=-1)
+
+    out = torch.matmul(attn, v.view(b_new, -1, c))  # [B*H*K, W/K, C]
+
+    out = merge_splits_1d(out, h, num_splits=num_splits)  # [B, H, W, C]
+
+    # shift back
+    if with_shift:
+        out = torch.roll(out, shifts=shift_size_w, dims=2)
+
+    out = out.view(b, -1, c)
+
+    return out
+
+
+class SelfAttnPropagation(nn.Module):
+    """
+    flow propagation with self-attention on feature
+    query: feature0, key: feature0, value: flow
+    """
+
+    def __init__(self, in_channels,
+                 **kwargs,
+                 ):
+        super(SelfAttnPropagation, self).__init__()
+
+        self.q_proj = nn.Linear(in_channels, in_channels)
+        self.k_proj = nn.Linear(in_channels, in_channels)
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, feature0, flow,
+                local_window_attn=False,
+                local_window_radius=1,
+                **kwargs,
+                ):
+        # q, k: feature [B, C, H, W], v: flow [B, 2, H, W]
+        if local_window_attn:
+            return self.forward_local_window_attn(feature0, flow,
+                                                  local_window_radius=local_window_radius)
+
+        b, c, h, w = feature0.size()
+
+        query = feature0.view(b, c, h * w).permute(0, 2, 1)  # [B, H*W, C]
+
+        # a note: the ``correct'' implementation should be:
+        # ``query = self.q_proj(query), key = self.k_proj(query)''
+        # this problem is observed while cleaning up the code
+        # however, this doesn't affect the performance since the projection is a linear operation,
+        # thus the two projection matrices for key can be merged
+        # so I just leave it as is in order to not re-train all models :)
+        query = self.q_proj(query)  # [B, H*W, C]
+        key = self.k_proj(query)  # [B, H*W, C]
+
+        value = flow.view(b, flow.size(1), h * w).permute(0, 2, 1)  # [B, H*W, 2]
+
+        scores = torch.matmul(query, key.permute(0, 2, 1)) / (c ** 0.5)  # [B, H*W, H*W]
+        prob = torch.softmax(scores, dim=-1)
+
+        out = torch.matmul(prob, value)  # [B, H*W, 2]
+        out = out.view(b, h, w, value.size(-1)).permute(0, 3, 1, 2)  # [B, 2, H, W]
+
+        return out
+
+    def forward_local_window_attn(self, feature0, flow,
+                                  local_window_radius=1,
+                                  ):
+        assert flow.size(1) == 2 or flow.size(1) == 1  # flow or disparity or depth
+        assert local_window_radius > 0
+
+        b, c, h, w = feature0.size()
+
+        value_channel = flow.size(1)
+
+        feature0_reshape = self.q_proj(feature0.view(b, c, -1).permute(0, 2, 1)
+                                       ).reshape(b * h * w, 1, c)  # [B*H*W, 1, C]
+
+        kernel_size = 2 * local_window_radius + 1
+
+        feature0_proj = self.k_proj(feature0.view(b, c, -1).permute(0, 2, 1)).permute(0, 2, 1).reshape(b, c, h, w)
+
+        feature0_window = F.unfold(feature0_proj, kernel_size=kernel_size,
+                                   padding=local_window_radius)  # [B, C*(2R+1)^2), H*W]
+
+        feature0_window = feature0_window.view(b, c, kernel_size ** 2, h, w).permute(
+            0, 3, 4, 1, 2).reshape(b * h * w, c, kernel_size ** 2)  # [B*H*W, C, (2R+1)^2]
+
+        flow_window = F.unfold(flow, kernel_size=kernel_size,
+                               padding=local_window_radius)  # [B, 2*(2R+1)^2), H*W]
+
+        flow_window = flow_window.view(b, value_channel, kernel_size ** 2, h, w).permute(
+            0, 3, 4, 2, 1).reshape(b * h * w, kernel_size ** 2, value_channel)  # [B*H*W, (2R+1)^2, 2]
+
+        scores = torch.matmul(feature0_reshape, feature0_window) / (c ** 0.5)  # [B*H*W, 1, (2R+1)^2]
+
+        prob = torch.softmax(scores, dim=-1)
+
+        out = torch.matmul(prob, flow_window).view(b, h, w, value_channel
+                                                   ).permute(0, 3, 1, 2).contiguous()  # [B, 2, H, W]
+
+        return out

unimatch/backbone.py

@@ -0,0 +1,117 @@
+import torch.nn as nn
+
+from .trident_conv import MultiScaleTridentConv
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_layer=nn.InstanceNorm2d, stride=1, dilation=1,
+                 ):
+        super(ResidualBlock, self).__init__()
+
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3,
+                               dilation=dilation, padding=dilation, stride=stride, bias=False)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                               dilation=dilation, padding=dilation, bias=False)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.norm1 = norm_layer(planes)
+        self.norm2 = norm_layer(planes)
+        if not stride == 1 or in_planes != planes:
+            self.norm3 = norm_layer(planes)
+
+        if stride == 1 and in_planes == planes:
+            self.downsample = None
+        else:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x + y)
+
+
+class CNNEncoder(nn.Module):
+    def __init__(self, output_dim=128,
+                 norm_layer=nn.InstanceNorm2d,
+                 num_output_scales=1,
+                 **kwargs,
+                 ):
+        super(CNNEncoder, self).__init__()
+        self.num_branch = num_output_scales
+
+        feature_dims = [64, 96, 128]
+
+        self.conv1 = nn.Conv2d(3, feature_dims[0], kernel_size=7, stride=2, padding=3, bias=False)  # 1/2
+        self.norm1 = norm_layer(feature_dims[0])
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = feature_dims[0]
+        self.layer1 = self._make_layer(feature_dims[0], stride=1, norm_layer=norm_layer)  # 1/2
+        self.layer2 = self._make_layer(feature_dims[1], stride=2, norm_layer=norm_layer)  # 1/4
+
+        # highest resolution 1/4 or 1/8
+        stride = 2 if num_output_scales == 1 else 1
+        self.layer3 = self._make_layer(feature_dims[2], stride=stride,
+                                       norm_layer=norm_layer,
+                                       )  # 1/4 or 1/8
+
+        self.conv2 = nn.Conv2d(feature_dims[2], output_dim, 1, 1, 0)
+
+        if self.num_branch > 1:
+            if self.num_branch == 4:
+                strides = (1, 2, 4, 8)
+            elif self.num_branch == 3:
+                strides = (1, 2, 4)
+            elif self.num_branch == 2:
+                strides = (1, 2)
+            else:
+                raise ValueError
+
+            self.trident_conv = MultiScaleTridentConv(output_dim, output_dim,
+                                                      kernel_size=3,
+                                                      strides=strides,
+                                                      paddings=1,
+                                                      num_branch=self.num_branch,
+                                                      )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1, dilation=1, norm_layer=nn.InstanceNorm2d):
+        layer1 = ResidualBlock(self.in_planes, dim, norm_layer=norm_layer, stride=stride, dilation=dilation)
+        layer2 = ResidualBlock(dim, dim, norm_layer=norm_layer, stride=1, dilation=dilation)
+
+        layers = (layer1, layer2)
+
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)  # 1/2
+        x = self.layer2(x)  # 1/4
+        x = self.layer3(x)  # 1/8 or 1/4
+
+        x = self.conv2(x)
+
+        if self.num_branch > 1:
+            out = self.trident_conv([x] * self.num_branch)  # high to low res
+        else:
+            out = [x]
+
+        return out

unimatch/geometry.py

@@ -0,0 +1,195 @@
+import torch
+import torch.nn.functional as F
+
+
+def coords_grid(b, h, w, homogeneous=False, device=None):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w))  # [H, W]
+
+    stacks = [x, y]
+
+    if homogeneous:
+        ones = torch.ones_like(x)  # [H, W]
+        stacks.append(ones)
+
+    grid = torch.stack(stacks, dim=0).float()  # [2, H, W] or [3, H, W]
+
+    grid = grid[None].repeat(b, 1, 1, 1)  # [B, 2, H, W] or [B, 3, H, W]
+
+    if device is not None:
+        grid = grid.to(device)
+
+    return grid
+
+
+def generate_window_grid(h_min, h_max, w_min, w_max, len_h, len_w, device=None):
+    assert device is not None
+
+    x, y = torch.meshgrid([torch.linspace(w_min, w_max, len_w, device=device),
+                           torch.linspace(h_min, h_max, len_h, device=device)],
+                          )
+    grid = torch.stack((x, y), -1).transpose(0, 1).float()  # [H, W, 2]
+
+    return grid
+
+
+def normalize_coords(coords, h, w):
+    # coords: [B, H, W, 2]
+    c = torch.Tensor([(w - 1) / 2., (h - 1) / 2.]).float().to(coords.device)
+    return (coords - c) / c  # [-1, 1]
+
+
+def bilinear_sample(img, sample_coords, mode='bilinear', padding_mode='zeros', return_mask=False):
+    # img: [B, C, H, W]
+    # sample_coords: [B, 2, H, W] in image scale
+    if sample_coords.size(1) != 2:  # [B, H, W, 2]
+        sample_coords = sample_coords.permute(0, 3, 1, 2)
+
+    b, _, h, w = sample_coords.shape
+
+    # Normalize to [-1, 1]
+    x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1
+    y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1
+
+    grid = torch.stack([x_grid, y_grid], dim=-1)  # [B, H, W, 2]
+
+    img = F.grid_sample(img, grid, mode=mode, padding_mode=padding_mode, align_corners=True)
+
+    if return_mask:
+        mask = (x_grid >= -1) & (y_grid >= -1) & (x_grid <= 1) & (y_grid <= 1)  # [B, H, W]
+
+        return img, mask
+
+    return img
+
+
+def flow_warp(feature, flow, mask=False, padding_mode='zeros'):
+    b, c, h, w = feature.size()
+    assert flow.size(1) == 2
+
+    grid = coords_grid(b, h, w).to(flow.device) + flow  # [B, 2, H, W]
+
+    return bilinear_sample(feature, grid, padding_mode=padding_mode,
+                           return_mask=mask)
+
+
+def forward_backward_consistency_check(fwd_flow, bwd_flow,
+                                       alpha=0.01,
+                                       beta=0.5
+                                       ):
+    # fwd_flow, bwd_flow: [B, 2, H, W]
+    # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837)
+    assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
+    assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
+    flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1)  # [B, H, W]
+
+    warped_bwd_flow = flow_warp(bwd_flow, fwd_flow)  # [B, 2, H, W]
+    warped_fwd_flow = flow_warp(fwd_flow, bwd_flow)  # [B, 2, H, W]
+
+    diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1)  # [B, H, W]
+    diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)
+
+    threshold = alpha * flow_mag + beta
+
+    fwd_occ = (diff_fwd > threshold).float()  # [B, H, W]
+    bwd_occ = (diff_bwd > threshold).float()
+
+    return fwd_occ, bwd_occ
+
+
+def back_project(depth, intrinsics):
+    # Back project 2D pixel coords to 3D points
+    # depth: [B, H, W]
+    # intrinsics: [B, 3, 3]
+    b, h, w = depth.shape
+    grid = coords_grid(b, h, w, homogeneous=True, device=depth.device)  # [B, 3, H, W]
+
+    intrinsics_inv = torch.inverse(intrinsics)  # [B, 3, 3]
+
+    points = intrinsics_inv.bmm(grid.view(b, 3, -1)).view(b, 3, h, w) * depth.unsqueeze(1)  # [B, 3, H, W]
+
+    return points
+
+
+def camera_transform(points_ref, extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None):
+    # Transform 3D points from reference camera to target camera
+    # points_ref: [B, 3, H, W]
+    # extrinsics_ref: [B, 4, 4]
+    # extrinsics_tgt: [B, 4, 4]
+    # extrinsics_rel: [B, 4, 4], relative pose transform
+    b, _, h, w = points_ref.shape
+
+    if extrinsics_rel is None:
+        extrinsics_rel = torch.bmm(extrinsics_tgt, torch.inverse(extrinsics_ref))  # [B, 4, 4]
+
+    points_tgt = torch.bmm(extrinsics_rel[:, :3, :3],
+                           points_ref.view(b, 3, -1)) + extrinsics_rel[:, :3, -1:]  # [B, 3, H*W]
+
+    points_tgt = points_tgt.view(b, 3, h, w)  # [B, 3, H, W]
+
+    return points_tgt
+
+
+def reproject(points_tgt, intrinsics, return_mask=False):
+    # reproject to target view
+    # points_tgt: [B, 3, H, W]
+    # intrinsics: [B, 3, 3]
+
+    b, _, h, w = points_tgt.shape
+
+    proj_points = torch.bmm(intrinsics, points_tgt.view(b, 3, -1)).view(b, 3, h, w)  # [B, 3, H, W]
+
+    X = proj_points[:, 0]
+    Y = proj_points[:, 1]
+    Z = proj_points[:, 2].clamp(min=1e-3)
+
+    pixel_coords = torch.stack([X / Z, Y / Z], dim=1).view(b, 2, h, w)  # [B, 2, H, W] in image scale
+
+    if return_mask:
+        # valid mask in pixel space
+        mask = (pixel_coords[:, 0] >= 0) & (pixel_coords[:, 0] <= (w - 1)) & (
+                pixel_coords[:, 1] >= 0) & (pixel_coords[:, 1] <= (h - 1))  # [B, H, W]
+
+        return pixel_coords, mask
+
+    return pixel_coords
+
+
+def reproject_coords(depth_ref, intrinsics, extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None,
+                     return_mask=False):
+    # Compute reprojection sample coords
+    points_ref = back_project(depth_ref, intrinsics)  # [B, 3, H, W]
+    points_tgt = camera_transform(points_ref, extrinsics_ref, extrinsics_tgt, extrinsics_rel=extrinsics_rel)
+
+    if return_mask:
+        reproj_coords, mask = reproject(points_tgt, intrinsics,
+                                        return_mask=return_mask)  # [B, 2, H, W] in image scale
+
+        return reproj_coords, mask
+
+    reproj_coords = reproject(points_tgt, intrinsics,
+                              return_mask=return_mask)  # [B, 2, H, W] in image scale
+
+    return reproj_coords
+
+
+def compute_flow_with_depth_pose(depth_ref, intrinsics,
+                                 extrinsics_ref=None, extrinsics_tgt=None, extrinsics_rel=None,
+                                 return_mask=False):
+    b, h, w = depth_ref.shape
+    coords_init = coords_grid(b, h, w, device=depth_ref.device)  # [B, 2, H, W]
+
+    if return_mask:
+        reproj_coords, mask = reproject_coords(depth_ref, intrinsics, extrinsics_ref, extrinsics_tgt,
+                                               extrinsics_rel=extrinsics_rel,
+                                               return_mask=return_mask)  # [B, 2, H, W]
+        rigid_flow = reproj_coords - coords_init
+
+        return rigid_flow, mask
+
+    reproj_coords = reproject_coords(depth_ref, intrinsics, extrinsics_ref, extrinsics_tgt,
+                                     extrinsics_rel=extrinsics_rel,
+                                     return_mask=return_mask)  # [B, 2, H, W]
+
+    rigid_flow = reproj_coords - coords_init
+
+    return rigid_flow

unimatch/matching.py

@@ -0,0 +1,279 @@
+import torch
+import torch.nn.functional as F
+
+from .geometry import coords_grid, generate_window_grid, normalize_coords
+
+
+def global_correlation_softmax(feature0, feature1,
+                               pred_bidir_flow=False,
+                               ):
+    # global correlation
+    b, c, h, w = feature0.shape
+    feature0 = feature0.view(b, c, -1).permute(0, 2, 1)  # [B, H*W, C]
+    feature1 = feature1.view(b, c, -1)  # [B, C, H*W]
+
+    correlation = torch.matmul(feature0, feature1).view(b, h, w, h, w) / (c ** 0.5)  # [B, H, W, H, W]
+
+    # flow from softmax
+    init_grid = coords_grid(b, h, w).to(correlation.device)  # [B, 2, H, W]
+    grid = init_grid.view(b, 2, -1).permute(0, 2, 1)  # [B, H*W, 2]
+
+    correlation = correlation.view(b, h * w, h * w)  # [B, H*W, H*W]
+
+    if pred_bidir_flow:
+        correlation = torch.cat((correlation, correlation.permute(0, 2, 1)), dim=0)  # [2*B, H*W, H*W]
+        init_grid = init_grid.repeat(2, 1, 1, 1)  # [2*B, 2, H, W]
+        grid = grid.repeat(2, 1, 1)  # [2*B, H*W, 2]
+        b = b * 2
+
+    prob = F.softmax(correlation, dim=-1)  # [B, H*W, H*W]
+
+    correspondence = torch.matmul(prob, grid).view(b, h, w, 2).permute(0, 3, 1, 2)  # [B, 2, H, W]
+
+    # when predicting bidirectional flow, flow is the concatenation of forward flow and backward flow
+    flow = correspondence - init_grid
+
+    return flow, prob
+
+
+def local_correlation_softmax(feature0, feature1, local_radius,
+                              padding_mode='zeros',
+                              ):
+    b, c, h, w = feature0.size()
+    coords_init = coords_grid(b, h, w).to(feature0.device)  # [B, 2, H, W]
+    coords = coords_init.view(b, 2, -1).permute(0, 2, 1)  # [B, H*W, 2]
+
+    local_h = 2 * local_radius + 1
+    local_w = 2 * local_radius + 1
+
+    window_grid = generate_window_grid(-local_radius, local_radius,
+                                       -local_radius, local_radius,
+                                       local_h, local_w, device=feature0.device)  # [2R+1, 2R+1, 2]
+    window_grid = window_grid.reshape(-1, 2).repeat(b, 1, 1, 1)  # [B, 1, (2R+1)^2, 2]
+    sample_coords = coords.unsqueeze(-2) + window_grid  # [B, H*W, (2R+1)^2, 2]
+
+    sample_coords_softmax = sample_coords
+
+    # exclude coords that are out of image space
+    valid_x = (sample_coords[:, :, :, 0] >= 0) & (sample_coords[:, :, :, 0] < w)  # [B, H*W, (2R+1)^2]
+    valid_y = (sample_coords[:, :, :, 1] >= 0) & (sample_coords[:, :, :, 1] < h)  # [B, H*W, (2R+1)^2]
+
+    valid = valid_x & valid_y  # [B, H*W, (2R+1)^2], used to mask out invalid values when softmax
+
+    # normalize coordinates to [-1, 1]
+    sample_coords_norm = normalize_coords(sample_coords, h, w)  # [-1, 1]
+    window_feature = F.grid_sample(feature1, sample_coords_norm,
+                                   padding_mode=padding_mode, align_corners=True
+                                   ).permute(0, 2, 1, 3)  # [B, H*W, C, (2R+1)^2]
+    feature0_view = feature0.permute(0, 2, 3, 1).view(b, h * w, 1, c)  # [B, H*W, 1, C]
+
+    corr = torch.matmul(feature0_view, window_feature).view(b, h * w, -1) / (c ** 0.5)  # [B, H*W, (2R+1)^2]
+
+    # mask invalid locations
+    corr[~valid] = -1e9
+
+    prob = F.softmax(corr, -1)  # [B, H*W, (2R+1)^2]
+
+    correspondence = torch.matmul(prob.unsqueeze(-2), sample_coords_softmax).squeeze(-2).view(
+        b, h, w, 2).permute(0, 3, 1, 2)  # [B, 2, H, W]
+
+    flow = correspondence - coords_init
+    match_prob = prob
+
+    return flow, match_prob
+
+
+def local_correlation_with_flow(feature0, feature1,
+                                flow,
+                                local_radius,
+                                padding_mode='zeros',
+                                dilation=1,
+                                ):
+    b, c, h, w = feature0.size()
+    coords_init = coords_grid(b, h, w).to(feature0.device)  # [B, 2, H, W]
+    coords = coords_init.view(b, 2, -1).permute(0, 2, 1)  # [B, H*W, 2]
+
+    local_h = 2 * local_radius + 1
+    local_w = 2 * local_radius + 1
+
+    window_grid = generate_window_grid(-local_radius, local_radius,
+                                       -local_radius, local_radius,
+                                       local_h, local_w, device=feature0.device)  # [2R+1, 2R+1, 2]
+    window_grid = window_grid.reshape(-1, 2).repeat(b, 1, 1, 1)  # [B, 1, (2R+1)^2, 2]
+    sample_coords = coords.unsqueeze(-2) + window_grid * dilation  # [B, H*W, (2R+1)^2, 2]
+
+    # flow can be zero when using features after transformer
+    if not isinstance(flow, float):
+        sample_coords = sample_coords + flow.view(
+            b, 2, -1).permute(0, 2, 1).unsqueeze(-2)  # [B, H*W, (2R+1)^2, 2]
+    else:
+        assert flow == 0.
+
+    # normalize coordinates to [-1, 1]
+    sample_coords_norm = normalize_coords(sample_coords, h, w)  # [-1, 1]
+    window_feature = F.grid_sample(feature1, sample_coords_norm,
+                                   padding_mode=padding_mode, align_corners=True
+                                   ).permute(0, 2, 1, 3)  # [B, H*W, C, (2R+1)^2]
+    feature0_view = feature0.permute(0, 2, 3, 1).view(b, h * w, 1, c)  # [B, H*W, 1, C]
+
+    corr = torch.matmul(feature0_view, window_feature).view(b, h * w, -1) / (c ** 0.5)  # [B, H*W, (2R+1)^2]
+
+    corr = corr.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous()  # [B, (2R+1)^2, H, W]
+
+    return corr
+
+
+def global_correlation_softmax_stereo(feature0, feature1,
+                                      ):
+    # global correlation on horizontal direction
+    b, c, h, w = feature0.shape
+
+    x_grid = torch.linspace(0, w - 1, w, device=feature0.device)  # [W]
+
+    feature0 = feature0.permute(0, 2, 3, 1)  # [B, H, W, C]
+    feature1 = feature1.permute(0, 2, 1, 3)  # [B, H, C, W]
+
+    correlation = torch.matmul(feature0, feature1) / (c ** 0.5)  # [B, H, W, W]
+
+    # mask subsequent positions to make disparity positive
+    mask = torch.triu(torch.ones((w, w)), diagonal=1).type_as(feature0)  # [W, W]
+    valid_mask = (mask == 0).unsqueeze(0).unsqueeze(0).repeat(b, h, 1, 1)  # [B, H, W, W]
+
+    correlation[~valid_mask] = -1e9
+
+    prob = F.softmax(correlation, dim=-1)  # [B, H, W, W]
+
+    correspondence = (x_grid.view(1, 1, 1, w) * prob).sum(-1)  # [B, H, W]
+
+    # NOTE: unlike flow, disparity is typically positive
+    disparity = x_grid.view(1, 1, w).repeat(b, h, 1) - correspondence  # [B, H, W]
+
+    return disparity.unsqueeze(1), prob  # feature resolution
+
+
+def local_correlation_softmax_stereo(feature0, feature1, local_radius,
+                                     ):
+    b, c, h, w = feature0.size()
+    coords_init = coords_grid(b, h, w).to(feature0.device)  # [B, 2, H, W]
+    coords = coords_init.view(b, 2, -1).permute(0, 2, 1).contiguous()  # [B, H*W, 2]
+
+    local_h = 1
+    local_w = 2 * local_radius + 1
+
+    window_grid = generate_window_grid(0, 0,
+                                       -local_radius, local_radius,
+                                       local_h, local_w, device=feature0.device)  # [1, 2R+1, 2]
+    window_grid = window_grid.reshape(-1, 2).repeat(b, 1, 1, 1)  # [B, 1, (2R+1), 2]
+    sample_coords = coords.unsqueeze(-2) + window_grid  # [B, H*W, (2R+1), 2]
+
+    sample_coords_softmax = sample_coords
+
+    # exclude coords that are out of image space
+    valid_x = (sample_coords[:, :, :, 0] >= 0) & (sample_coords[:, :, :, 0] < w)  # [B, H*W, (2R+1)^2]
+    valid_y = (sample_coords[:, :, :, 1] >= 0) & (sample_coords[:, :, :, 1] < h)  # [B, H*W, (2R+1)^2]
+
+    valid = valid_x & valid_y  # [B, H*W, (2R+1)^2], used to mask out invalid values when softmax
+
+    # normalize coordinates to [-1, 1]
+    sample_coords_norm = normalize_coords(sample_coords, h, w)  # [-1, 1]
+    window_feature = F.grid_sample(feature1, sample_coords_norm,
+                                   padding_mode='zeros', align_corners=True
+                                   ).permute(0, 2, 1, 3)  # [B, H*W, C, (2R+1)]
+    feature0_view = feature0.permute(0, 2, 3, 1).contiguous().view(b, h * w, 1, c)  # [B, H*W, 1, C]
+
+    corr = torch.matmul(feature0_view, window_feature).view(b, h * w, -1) / (c ** 0.5)  # [B, H*W, (2R+1)]
+
+    # mask invalid locations
+    corr[~valid] = -1e9
+
+    prob = F.softmax(corr, -1)  # [B, H*W, (2R+1)]
+
+    correspondence = torch.matmul(prob.unsqueeze(-2),
+                                  sample_coords_softmax).squeeze(-2).view(
+        b, h, w, 2).permute(0, 3, 1, 2).contiguous()  # [B, 2, H, W]
+
+    flow = correspondence - coords_init  # flow at feature resolution
+    match_prob = prob
+
+    flow_x = -flow[:, :1]  # [B, 1, H, W]
+
+    return flow_x, match_prob
+
+
+def correlation_softmax_depth(feature0, feature1,
+                              intrinsics,
+                              pose,
+                              depth_candidates,
+                              depth_from_argmax=False,
+                              pred_bidir_depth=False,
+                              ):
+    b, c, h, w = feature0.size()
+    assert depth_candidates.dim() == 4  # [B, D, H, W]
+    scale_factor = c ** 0.5
+
+    if pred_bidir_depth:
+        feature0, feature1 = torch.cat((feature0, feature1), dim=0), torch.cat((feature1, feature0), dim=0)
+        intrinsics = intrinsics.repeat(2, 1, 1)
+        pose = torch.cat((pose, torch.inverse(pose)), dim=0)
+        depth_candidates = depth_candidates.repeat(2, 1, 1, 1)
+
+    # depth candidates are actually inverse depth
+    warped_feature1 = warp_with_pose_depth_candidates(feature1, intrinsics, pose,
+                                                      1. / depth_candidates,
+                                                      )  # [B, C, D, H, W]
+
+    correlation = (feature0.unsqueeze(2) * warped_feature1).sum(1) / scale_factor  # [B, D, H, W]
+
+    match_prob = F.softmax(correlation, dim=1)  # [B, D, H, W]
+
+    # for cross-task transfer (flow -> depth), extract depth with argmax at test time
+    if depth_from_argmax:
+        index = torch.argmax(match_prob, dim=1, keepdim=True)
+        depth = torch.gather(depth_candidates, dim=1, index=index)
+    else:
+        depth = (match_prob * depth_candidates).sum(dim=1, keepdim=True)  # [B, 1, H, W]
+
+    return depth, match_prob
+
+
+def warp_with_pose_depth_candidates(feature1, intrinsics, pose, depth,
+                                    clamp_min_depth=1e-3,
+                                    ):
+    """
+    feature1: [B, C, H, W]
+    intrinsics: [B, 3, 3]
+    pose: [B, 4, 4]
+    depth: [B, D, H, W]
+    """
+
+    assert intrinsics.size(1) == intrinsics.size(2) == 3
+    assert pose.size(1) == pose.size(2) == 4
+    assert depth.dim() == 4
+
+    b, d, h, w = depth.size()
+    c = feature1.size(1)
+
+    with torch.no_grad():
+        # pixel coordinates
+        grid = coords_grid(b, h, w, homogeneous=True, device=depth.device)  # [B, 3, H, W]
+        # back project to 3D and transform viewpoint
+        points = torch.inverse(intrinsics).bmm(grid.view(b, 3, -1))  # [B, 3, H*W]
+        points = torch.bmm(pose[:, :3, :3], points).unsqueeze(2).repeat(
+            1, 1, d, 1) * depth.view(b, 1, d, h * w)  # [B, 3, D, H*W]
+        points = points + pose[:, :3, -1:].unsqueeze(-1)  # [B, 3, D, H*W]
+        # reproject to 2D image plane
+        points = torch.bmm(intrinsics, points.view(b, 3, -1)).view(b, 3, d, h * w)  # [B, 3, D, H*W]
+        pixel_coords = points[:, :2] / points[:, -1:].clamp(min=clamp_min_depth)  # [B, 2, D, H*W]
+
+        # normalize to [-1, 1]
+        x_grid = 2 * pixel_coords[:, 0] / (w - 1) - 1
+        y_grid = 2 * pixel_coords[:, 1] / (h - 1) - 1
+
+        grid = torch.stack([x_grid, y_grid], dim=-1)  # [B, D, H*W, 2]
+
+    # sample features
+    warped_feature = F.grid_sample(feature1, grid.view(b, d * h, w, 2), mode='bilinear',
+                                   padding_mode='zeros',
+                                   align_corners=True).view(b, c, d, h, w)  # [B, C, D, H, W]
+
+    return warped_feature

unimatch/position.py

@@ -0,0 +1,46 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# https://github.com/facebookresearch/detr/blob/main/models/position_encoding.py
+
+import torch
+import torch.nn as nn
+import math
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=True, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, x):
+        # x = tensor_list.tensors  # [B, C, H, W]
+        # mask = tensor_list.mask  # [B, H, W], input with padding, valid as 0
+        b, c, h, w = x.size()
+        mask = torch.ones((b, h, w), device=x.device)  # [B, H, W]
+        y_embed = mask.cumsum(1, dtype=torch.float32)
+        x_embed = mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        return pos

unimatch/reg_refine.py

@@ -0,0 +1,119 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FlowHead(nn.Module):
+    def __init__(self, input_dim=128, hidden_dim=256,
+                 out_dim=2,
+                 ):
+        super(FlowHead, self).__init__()
+
+        self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
+        self.conv2 = nn.Conv2d(hidden_dim, out_dim, 3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        out = self.conv2(self.relu(self.conv1(x)))
+
+        return out
+
+
+class SepConvGRU(nn.Module):
+    def __init__(self, hidden_dim=128, input_dim=192 + 128,
+                 kernel_size=5,
+                 ):
+        padding = (kernel_size - 1) // 2
+
+        super(SepConvGRU, self).__init__()
+        self.convz1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, kernel_size), padding=(0, padding))
+        self.convr1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, kernel_size), padding=(0, padding))
+        self.convq1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, kernel_size), padding=(0, padding))
+
+        self.convz2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (kernel_size, 1), padding=(padding, 0))
+        self.convr2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (kernel_size, 1), padding=(padding, 0))
+        self.convq2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (kernel_size, 1), padding=(padding, 0))
+
+    def forward(self, h, x):
+        # horizontal
+        hx = torch.cat([h, x], dim=1)
+        z = torch.sigmoid(self.convz1(hx))
+        r = torch.sigmoid(self.convr1(hx))
+        q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1)))
+        h = (1 - z) * h + z * q
+
+        # vertical
+        hx = torch.cat([h, x], dim=1)
+        z = torch.sigmoid(self.convz2(hx))
+        r = torch.sigmoid(self.convr2(hx))
+        q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1)))
+        h = (1 - z) * h + z * q
+
+        return h
+
+
+class BasicMotionEncoder(nn.Module):
+    def __init__(self, corr_channels=324,
+                 flow_channels=2,
+                 ):
+        super(BasicMotionEncoder, self).__init__()
+
+        self.convc1 = nn.Conv2d(corr_channels, 256, 1, padding=0)
+        self.convc2 = nn.Conv2d(256, 192, 3, padding=1)
+        self.convf1 = nn.Conv2d(flow_channels, 128, 7, padding=3)
+        self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
+        self.conv = nn.Conv2d(64 + 192, 128 - flow_channels, 3, padding=1)
+
+    def forward(self, flow, corr):
+        cor = F.relu(self.convc1(corr))
+        cor = F.relu(self.convc2(cor))
+        flo = F.relu(self.convf1(flow))
+        flo = F.relu(self.convf2(flo))
+
+        cor_flo = torch.cat([cor, flo], dim=1)
+        out = F.relu(self.conv(cor_flo))
+        return torch.cat([out, flow], dim=1)
+
+
+class BasicUpdateBlock(nn.Module):
+    def __init__(self, corr_channels=324,
+                 hidden_dim=128,
+                 context_dim=128,
+                 downsample_factor=8,
+                 flow_dim=2,
+                 bilinear_up=False,
+                 ):
+        super(BasicUpdateBlock, self).__init__()
+
+        self.encoder = BasicMotionEncoder(corr_channels=corr_channels,
+                                          flow_channels=flow_dim,
+                                          )
+
+        self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=context_dim + hidden_dim)
+
+        self.flow_head = FlowHead(hidden_dim, hidden_dim=256,
+                                  out_dim=flow_dim,
+                                  )
+
+        if bilinear_up:
+            self.mask = None
+        else:
+            self.mask = nn.Sequential(
+                nn.Conv2d(hidden_dim, 256, 3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(256, downsample_factor ** 2 * 9, 1, padding=0))
+
+    def forward(self, net, inp, corr, flow):
+        motion_features = self.encoder(flow, corr)
+
+        inp = torch.cat([inp, motion_features], dim=1)
+
+        net = self.gru(net, inp)
+        delta_flow = self.flow_head(net)
+
+        if self.mask is not None:
+            mask = self.mask(net)
+        else:
+            mask = None
+
+        return net, mask, delta_flow

unimatch/transformer.py

@@ -0,0 +1,294 @@
+import torch
+import torch.nn as nn
+
+from .attention import (single_head_full_attention, single_head_split_window_attention,
+                        single_head_full_attention_1d, single_head_split_window_attention_1d)
+from .utils import generate_shift_window_attn_mask, generate_shift_window_attn_mask_1d
+
+
+class TransformerLayer(nn.Module):
+    def __init__(self,
+                 d_model=128,
+                 nhead=1,
+                 no_ffn=False,
+                 ffn_dim_expansion=4,
+                 ):
+        super(TransformerLayer, self).__init__()
+
+        self.dim = d_model
+        self.nhead = nhead
+        self.no_ffn = no_ffn
+
+        # multi-head attention
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+
+        self.merge = nn.Linear(d_model, d_model, bias=False)
+
+        self.norm1 = nn.LayerNorm(d_model)
+
+        # no ffn after self-attn, with ffn after cross-attn
+        if not self.no_ffn:
+            in_channels = d_model * 2
+            self.mlp = nn.Sequential(
+                nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False),
+                nn.GELU(),
+                nn.Linear(in_channels * ffn_dim_expansion, d_model, bias=False),
+            )
+
+            self.norm2 = nn.LayerNorm(d_model)
+
+    def forward(self, source, target,
+                height=None,
+                width=None,
+                shifted_window_attn_mask=None,
+                shifted_window_attn_mask_1d=None,
+                attn_type='swin',
+                with_shift=False,
+                attn_num_splits=None,
+                ):
+        # source, target: [B, L, C]
+        query, key, value = source, target, target
+
+        # for stereo: 2d attn in self-attn, 1d attn in cross-attn
+        is_self_attn = (query - key).abs().max() < 1e-6
+
+        # single-head attention
+        query = self.q_proj(query)  # [B, L, C]
+        key = self.k_proj(key)  # [B, L, C]
+        value = self.v_proj(value)  # [B, L, C]
+
+        if attn_type == 'swin' and attn_num_splits > 1:  # self, cross-attn: both swin 2d
+            if self.nhead > 1:
+                # we observe that multihead attention slows down the speed and increases the memory consumption
+                # without bringing obvious performance gains and thus the implementation is removed
+                raise NotImplementedError
+            else:
+                message = single_head_split_window_attention(query, key, value,
+                                                             num_splits=attn_num_splits,
+                                                             with_shift=with_shift,
+                                                             h=height,
+                                                             w=width,
+                                                             attn_mask=shifted_window_attn_mask,
+                                                             )
+
+        elif attn_type == 'self_swin2d_cross_1d':  # self-attn: swin 2d, cross-attn: full 1d
+            if self.nhead > 1:
+                raise NotImplementedError
+            else:
+                if is_self_attn:
+                    if attn_num_splits > 1:
+                        message = single_head_split_window_attention(query, key, value,
+                                                                     num_splits=attn_num_splits,
+                                                                     with_shift=with_shift,
+                                                                     h=height,
+                                                                     w=width,
+                                                                     attn_mask=shifted_window_attn_mask,
+                                                                     )
+                    else:
+                        # full 2d attn
+                        message = single_head_full_attention(query, key, value)  # [N, L, C]
+
+                else:
+                    # cross attn 1d
+                    message = single_head_full_attention_1d(query, key, value,
+                                                            h=height,
+                                                            w=width,
+                                                            )
+
+        elif attn_type == 'self_swin2d_cross_swin1d':  # self-attn: swin 2d, cross-attn: swin 1d
+            if self.nhead > 1:
+                raise NotImplementedError
+            else:
+                if is_self_attn:
+                    if attn_num_splits > 1:
+                        # self attn shift window
+                        message = single_head_split_window_attention(query, key, value,
+                                                                     num_splits=attn_num_splits,
+                                                                     with_shift=with_shift,
+                                                                     h=height,
+                                                                     w=width,
+                                                                     attn_mask=shifted_window_attn_mask,
+                                                                     )
+                    else:
+                        # full 2d attn
+                        message = single_head_full_attention(query, key, value)  # [N, L, C]
+                else:
+                    if attn_num_splits > 1:
+                        assert shifted_window_attn_mask_1d is not None
+                        # cross attn 1d shift
+                        message = single_head_split_window_attention_1d(query, key, value,
+                                                                        num_splits=attn_num_splits,
+                                                                        with_shift=with_shift,
+                                                                        h=height,
+                                                                        w=width,
+                                                                        attn_mask=shifted_window_attn_mask_1d,
+                                                                        )
+                    else:
+                        message = single_head_full_attention_1d(query, key, value,
+                                                                h=height,
+                                                                w=width,
+                                                                )
+
+        else:
+            message = single_head_full_attention(query, key, value)  # [B, L, C]
+
+        message = self.merge(message)  # [B, L, C]
+        message = self.norm1(message)
+
+        if not self.no_ffn:
+            message = self.mlp(torch.cat([source, message], dim=-1))
+            message = self.norm2(message)
+
+        return source + message
+
+
+class TransformerBlock(nn.Module):
+    """self attention + cross attention + FFN"""
+
+    def __init__(self,
+                 d_model=128,
+                 nhead=1,
+                 ffn_dim_expansion=4,
+                 ):
+        super(TransformerBlock, self).__init__()
+
+        self.self_attn = TransformerLayer(d_model=d_model,
+                                          nhead=nhead,
+                                          no_ffn=True,
+                                          ffn_dim_expansion=ffn_dim_expansion,
+                                          )
+
+        self.cross_attn_ffn = TransformerLayer(d_model=d_model,
+                                               nhead=nhead,
+                                               ffn_dim_expansion=ffn_dim_expansion,
+                                               )
+
+    def forward(self, source, target,
+                height=None,
+                width=None,
+                shifted_window_attn_mask=None,
+                shifted_window_attn_mask_1d=None,
+                attn_type='swin',
+                with_shift=False,
+                attn_num_splits=None,
+                ):
+        # source, target: [B, L, C]
+
+        # self attention
+        source = self.self_attn(source, source,
+                                height=height,
+                                width=width,
+                                shifted_window_attn_mask=shifted_window_attn_mask,
+                                attn_type=attn_type,
+                                with_shift=with_shift,
+                                attn_num_splits=attn_num_splits,
+                                )
+
+        # cross attention and ffn
+        source = self.cross_attn_ffn(source, target,
+                                     height=height,
+                                     width=width,
+                                     shifted_window_attn_mask=shifted_window_attn_mask,
+                                     shifted_window_attn_mask_1d=shifted_window_attn_mask_1d,
+                                     attn_type=attn_type,
+                                     with_shift=with_shift,
+                                     attn_num_splits=attn_num_splits,
+                                     )
+
+        return source
+
+
+class FeatureTransformer(nn.Module):
+    def __init__(self,
+                 num_layers=6,
+                 d_model=128,
+                 nhead=1,
+                 ffn_dim_expansion=4,
+                 ):
+        super(FeatureTransformer, self).__init__()
+
+        self.d_model = d_model
+        self.nhead = nhead
+
+        self.layers = nn.ModuleList([
+            TransformerBlock(d_model=d_model,
+                             nhead=nhead,
+                             ffn_dim_expansion=ffn_dim_expansion,
+                             )
+            for i in range(num_layers)])
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, feature0, feature1,
+                attn_type='swin',
+                attn_num_splits=None,
+                **kwargs,
+                ):
+
+        b, c, h, w = feature0.shape
+        assert self.d_model == c
+
+        feature0 = feature0.flatten(-2).permute(0, 2, 1)  # [B, H*W, C]
+        feature1 = feature1.flatten(-2).permute(0, 2, 1)  # [B, H*W, C]
+
+        # 2d attention
+        if 'swin' in attn_type and attn_num_splits > 1:
+            # global and refine use different number of splits
+            window_size_h = h // attn_num_splits
+            window_size_w = w // attn_num_splits
+
+            # compute attn mask once
+            shifted_window_attn_mask = generate_shift_window_attn_mask(
+                input_resolution=(h, w),
+                window_size_h=window_size_h,
+                window_size_w=window_size_w,
+                shift_size_h=window_size_h // 2,
+                shift_size_w=window_size_w // 2,
+                device=feature0.device,
+            )  # [K*K, H/K*W/K, H/K*W/K]
+        else:
+            shifted_window_attn_mask = None
+
+        # 1d attention
+        if 'swin1d' in attn_type and attn_num_splits > 1:
+            window_size_w = w // attn_num_splits
+
+            # compute attn mask once
+            shifted_window_attn_mask_1d = generate_shift_window_attn_mask_1d(
+                input_w=w,
+                window_size_w=window_size_w,
+                shift_size_w=window_size_w // 2,
+                device=feature0.device,
+            )  # [K, W/K, W/K]
+        else:
+            shifted_window_attn_mask_1d = None
+
+        # concat feature0 and feature1 in batch dimension to compute in parallel
+        concat0 = torch.cat((feature0, feature1), dim=0)  # [2B, H*W, C]
+        concat1 = torch.cat((feature1, feature0), dim=0)  # [2B, H*W, C]
+
+        for i, layer in enumerate(self.layers):
+            concat0 = layer(concat0, concat1,
+                            height=h,
+                            width=w,
+                            attn_type=attn_type,
+                            with_shift='swin' in attn_type and attn_num_splits > 1 and i % 2 == 1,
+                            attn_num_splits=attn_num_splits,
+                            shifted_window_attn_mask=shifted_window_attn_mask,
+                            shifted_window_attn_mask_1d=shifted_window_attn_mask_1d,
+                            )
+
+            # update feature1
+            concat1 = torch.cat(concat0.chunk(chunks=2, dim=0)[::-1], dim=0)
+
+        feature0, feature1 = concat0.chunk(chunks=2, dim=0)  # [B, H*W, C]
+
+        # reshape back
+        feature0 = feature0.view(b, h, w, c).permute(0, 3, 1, 2).contiguous()  # [B, C, H, W]
+        feature1 = feature1.view(b, h, w, c).permute(0, 3, 1, 2).contiguous()  # [B, C, H, W]
+
+        return feature0, feature1

unimatch/trident_conv.py

@@ -0,0 +1,90 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# https://github.com/facebookresearch/detectron2/blob/main/projects/TridentNet/tridentnet/trident_conv.py
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.modules.utils import _pair
+
+
+class MultiScaleTridentConv(nn.Module):
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=1,
+            strides=1,
+            paddings=0,
+            dilations=1,
+            dilation=1,
+            groups=1,
+            num_branch=1,
+            test_branch_idx=-1,
+            bias=False,
+            norm=None,
+            activation=None,
+    ):
+        super(MultiScaleTridentConv, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.num_branch = num_branch
+        self.stride = _pair(stride)
+        self.groups = groups
+        self.with_bias = bias
+        self.dilation = dilation
+        if isinstance(paddings, int):
+            paddings = [paddings] * self.num_branch
+        if isinstance(dilations, int):
+            dilations = [dilations] * self.num_branch
+        if isinstance(strides, int):
+            strides = [strides] * self.num_branch
+        self.paddings = [_pair(padding) for padding in paddings]
+        self.dilations = [_pair(dilation) for dilation in dilations]
+        self.strides = [_pair(stride) for stride in strides]
+        self.test_branch_idx = test_branch_idx
+        self.norm = norm
+        self.activation = activation
+
+        assert len({self.num_branch, len(self.paddings), len(self.strides)}) == 1
+
+        self.weight = nn.Parameter(
+            torch.Tensor(out_channels, in_channels // groups, *self.kernel_size)
+        )
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_channels))
+        else:
+            self.bias = None
+
+        nn.init.kaiming_uniform_(self.weight, nonlinearity="relu")
+        if self.bias is not None:
+            nn.init.constant_(self.bias, 0)
+
+    def forward(self, inputs):
+        num_branch = self.num_branch if self.training or self.test_branch_idx == -1 else 1
+        assert len(inputs) == num_branch
+
+        if self.training or self.test_branch_idx == -1:
+            outputs = [
+                F.conv2d(input, self.weight, self.bias, stride, padding, self.dilation, self.groups)
+                for input, stride, padding in zip(inputs, self.strides, self.paddings)
+            ]
+        else:
+            outputs = [
+                F.conv2d(
+                    inputs[0],
+                    self.weight,
+                    self.bias,
+                    self.strides[self.test_branch_idx] if self.test_branch_idx == -1 else self.strides[-1],
+                    self.paddings[self.test_branch_idx] if self.test_branch_idx == -1 else self.paddings[-1],
+                    self.dilation,
+                    self.groups,
+                )
+            ]
+
+        if self.norm is not None:
+            outputs = [self.norm(x) for x in outputs]
+        if self.activation is not None:
+            outputs = [self.activation(x) for x in outputs]
+        return outputs

unimatch/unimatch.py

@@ -0,0 +1,367 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbone import CNNEncoder
+from .transformer import FeatureTransformer
+from .matching import (global_correlation_softmax, local_correlation_softmax, local_correlation_with_flow,
+                       global_correlation_softmax_stereo, local_correlation_softmax_stereo,
+                       correlation_softmax_depth)
+from .attention import SelfAttnPropagation
+from .geometry import flow_warp, compute_flow_with_depth_pose
+from .reg_refine import BasicUpdateBlock
+from .utils import normalize_img, feature_add_position, upsample_flow_with_mask
+
+
+class UniMatch(nn.Module):
+    def __init__(self,
+                 num_scales=1,
+                 feature_channels=128,
+                 upsample_factor=8,
+                 num_head=1,
+                 ffn_dim_expansion=4,
+                 num_transformer_layers=6,
+                 reg_refine=False,  # optional local regression refinement
+                 task='flow',
+                 ):
+        super(UniMatch, self).__init__()
+
+        self.feature_channels = feature_channels
+        self.num_scales = num_scales
+        self.upsample_factor = upsample_factor
+        self.reg_refine = reg_refine
+
+        # CNN
+        self.backbone = CNNEncoder(output_dim=feature_channels, num_output_scales=num_scales)
+
+        # Transformer
+        self.transformer = FeatureTransformer(num_layers=num_transformer_layers,
+                                              d_model=feature_channels,
+                                              nhead=num_head,
+                                              ffn_dim_expansion=ffn_dim_expansion,
+                                              )
+
+        # propagation with self-attn
+        self.feature_flow_attn = SelfAttnPropagation(in_channels=feature_channels)
+
+        if not self.reg_refine or task == 'depth':
+            # convex upsampling simiar to RAFT
+            # concat feature0 and low res flow as input
+            self.upsampler = nn.Sequential(nn.Conv2d(2 + feature_channels, 256, 3, 1, 1),
+                                           nn.ReLU(inplace=True),
+                                           nn.Conv2d(256, upsample_factor ** 2 * 9, 1, 1, 0))
+            # thus far, all the learnable parameters are task-agnostic
+
+        if reg_refine:
+            # optional task-specific local regression refinement
+            self.refine_proj = nn.Conv2d(128, 256, 1)
+            self.refine = BasicUpdateBlock(corr_channels=(2 * 4 + 1) ** 2,
+                                           downsample_factor=upsample_factor,
+                                           flow_dim=2 if task == 'flow' else 1,
+                                           bilinear_up=task == 'depth',
+                                           )
+
+    def extract_feature(self, img0, img1):
+        concat = torch.cat((img0, img1), dim=0)  # [2B, C, H, W]
+        features = self.backbone(concat)  # list of [2B, C, H, W], resolution from high to low
+
+        # reverse: resolution from low to high
+        features = features[::-1]
+
+        feature0, feature1 = [], []
+
+        for i in range(len(features)):
+            feature = features[i]
+            chunks = torch.chunk(feature, 2, 0)  # tuple
+            feature0.append(chunks[0])
+            feature1.append(chunks[1])
+
+        return feature0, feature1
+
+    def upsample_flow(self, flow, feature, bilinear=False, upsample_factor=8,
+                      is_depth=False):
+        if bilinear:
+            multiplier = 1 if is_depth else upsample_factor
+            up_flow = F.interpolate(flow, scale_factor=upsample_factor,
+                                    mode='bilinear', align_corners=True) * multiplier
+        else:
+            concat = torch.cat((flow, feature), dim=1)
+            mask = self.upsampler(concat)
+            up_flow = upsample_flow_with_mask(flow, mask, upsample_factor=self.upsample_factor,
+                                              is_depth=is_depth)
+
+        return up_flow
+
+    def forward(self, img0, img1,
+                attn_type=None,
+                attn_splits_list=None,
+                corr_radius_list=None,
+                prop_radius_list=None,
+                num_reg_refine=1,
+                pred_bidir_flow=False,
+                task='flow',
+                intrinsics=None,
+                pose=None,  # relative pose transform
+                min_depth=1. / 0.5,  # inverse depth range
+                max_depth=1. / 10,
+                num_depth_candidates=64,
+                depth_from_argmax=False,
+                pred_bidir_depth=False,
+                **kwargs,
+                ):
+
+        if pred_bidir_flow:
+            assert task == 'flow'
+
+        if task == 'depth':
+            assert self.num_scales == 1  # multi-scale depth model is not supported yet
+
+        results_dict = {}
+        flow_preds = []
+
+        if task == 'flow':
+            # stereo and depth tasks have normalized img in dataloader
+            img0, img1 = normalize_img(img0, img1)  # [B, 3, H, W]
+
+        # list of features, resolution low to high
+        feature0_list, feature1_list = self.extract_feature(img0, img1)  # list of features
+
+        flow = None
+
+        if task != 'depth':
+            assert len(attn_splits_list) == len(corr_radius_list) == len(prop_radius_list) == self.num_scales
+        else:
+            assert len(attn_splits_list) == len(prop_radius_list) == self.num_scales == 1
+
+        for scale_idx in range(self.num_scales):
+            feature0, feature1 = feature0_list[scale_idx], feature1_list[scale_idx]
+
+            if pred_bidir_flow and scale_idx > 0:
+                # predicting bidirectional flow with refinement
+                feature0, feature1 = torch.cat((feature0, feature1), dim=0), torch.cat((feature1, feature0), dim=0)
+
+            feature0_ori, feature1_ori = feature0, feature1
+
+            upsample_factor = self.upsample_factor * (2 ** (self.num_scales - 1 - scale_idx))
+
+            if task == 'depth':
+                # scale intrinsics
+                intrinsics_curr = intrinsics.clone()
+                intrinsics_curr[:, :2] = intrinsics_curr[:, :2] / upsample_factor
+
+            if scale_idx > 0:
+                assert task != 'depth'  # not supported for multi-scale depth model
+                flow = F.interpolate(flow, scale_factor=2, mode='bilinear', align_corners=True) * 2
+
+            if flow is not None:
+                assert task != 'depth'
+                flow = flow.detach()
+
+                if task == 'stereo':
+                    # construct flow vector for disparity
+                    # flow here is actually disparity
+                    zeros = torch.zeros_like(flow)  # [B, 1, H, W]
+                    # NOTE: reverse disp, disparity is positive
+                    displace = torch.cat((-flow, zeros), dim=1)  # [B, 2, H, W]
+                    feature1 = flow_warp(feature1, displace)  # [B, C, H, W]
+                elif task == 'flow':
+                    feature1 = flow_warp(feature1, flow)  # [B, C, H, W]
+                else:
+                    raise NotImplementedError
+
+            attn_splits = attn_splits_list[scale_idx]
+            if task != 'depth':
+                corr_radius = corr_radius_list[scale_idx]
+            prop_radius = prop_radius_list[scale_idx]
+
+            # add position to features
+            feature0, feature1 = feature_add_position(feature0, feature1, attn_splits, self.feature_channels)
+
+            # Transformer
+            feature0, feature1 = self.transformer(feature0, feature1,
+                                                  attn_type=attn_type,
+                                                  attn_num_splits=attn_splits,
+                                                  )
+
+            # correlation and softmax
+            if task == 'depth':
+                # first generate depth candidates
+                b, _, h, w = feature0.size()
+                depth_candidates = torch.linspace(min_depth, max_depth, num_depth_candidates).type_as(feature0)
+                depth_candidates = depth_candidates.view(1, num_depth_candidates, 1, 1).repeat(b, 1, h,
+                                                                                               w)  # [B, D, H, W]
+
+                flow_pred = correlation_softmax_depth(feature0, feature1,
+                                                      intrinsics_curr,
+                                                      pose,
+                                                      depth_candidates=depth_candidates,
+                                                      depth_from_argmax=depth_from_argmax,
+                                                      pred_bidir_depth=pred_bidir_depth,
+                                                      )[0]
+
+            else:
+                if corr_radius == -1:  # global matching
+                    if task == 'flow':
+                        flow_pred = global_correlation_softmax(feature0, feature1, pred_bidir_flow)[0]
+                    elif task == 'stereo':
+                        flow_pred = global_correlation_softmax_stereo(feature0, feature1)[0]
+                    else:
+                        raise NotImplementedError
+                else:  # local matching
+                    if task == 'flow':
+                        flow_pred = local_correlation_softmax(feature0, feature1, corr_radius)[0]
+                    elif task == 'stereo':
+                        flow_pred = local_correlation_softmax_stereo(feature0, feature1, corr_radius)[0]
+                    else:
+                        raise NotImplementedError
+
+            # flow or residual flow
+            flow = flow + flow_pred if flow is not None else flow_pred
+
+            if task == 'stereo':
+                flow = flow.clamp(min=0)  # positive disparity
+
+            # upsample to the original resolution for supervison at training time only
+            if self.training:
+                flow_bilinear = self.upsample_flow(flow, None, bilinear=True, upsample_factor=upsample_factor,
+                                                   is_depth=task == 'depth')
+                flow_preds.append(flow_bilinear)
+
+            # flow propagation with self-attn
+            if (pred_bidir_flow or pred_bidir_depth) and scale_idx == 0:
+                feature0 = torch.cat((feature0, feature1), dim=0)  # [2*B, C, H, W] for propagation
+
+            flow = self.feature_flow_attn(feature0, flow.detach(),
+                                          local_window_attn=prop_radius > 0,
+                                          local_window_radius=prop_radius,
+                                          )
+
+            # bilinear exclude the last one
+            if self.training and scale_idx < self.num_scales - 1:
+                flow_up = self.upsample_flow(flow, feature0, bilinear=True,
+                                             upsample_factor=upsample_factor,
+                                             is_depth=task == 'depth')
+                flow_preds.append(flow_up)
+
+            if scale_idx == self.num_scales - 1:
+                if not self.reg_refine:
+                    # upsample to the original image resolution
+
+                    if task == 'stereo':
+                        flow_pad = torch.cat((-flow, torch.zeros_like(flow)), dim=1)  # [B, 2, H, W]
+                        flow_up_pad = self.upsample_flow(flow_pad, feature0)
+                        flow_up = -flow_up_pad[:, :1]  # [B, 1, H, W]
+                    elif task == 'depth':
+                        depth_pad = torch.cat((flow, torch.zeros_like(flow)), dim=1)  # [B, 2, H, W]
+                        depth_up_pad = self.upsample_flow(depth_pad, feature0,
+                                                          is_depth=True).clamp(min=min_depth, max=max_depth)
+                        flow_up = depth_up_pad[:, :1]  # [B, 1, H, W]
+                    else:
+                        flow_up = self.upsample_flow(flow, feature0)
+
+                    flow_preds.append(flow_up)
+                else:
+                    # task-specific local regression refinement
+                    # supervise current flow
+                    if self.training:
+                        flow_up = self.upsample_flow(flow, feature0, bilinear=True,
+                                                     upsample_factor=upsample_factor,
+                                                     is_depth=task == 'depth')
+                        flow_preds.append(flow_up)
+
+                    assert num_reg_refine > 0
+                    for refine_iter_idx in range(num_reg_refine):
+                        flow = flow.detach()
+
+                        if task == 'stereo':
+                            zeros = torch.zeros_like(flow)  # [B, 1, H, W]
+                            # NOTE: reverse disp, disparity is positive
+                            displace = torch.cat((-flow, zeros), dim=1)  # [B, 2, H, W]
+                            correlation = local_correlation_with_flow(
+                                feature0_ori,
+                                feature1_ori,
+                                flow=displace,
+                                local_radius=4,
+                            )  # [B, (2R+1)^2, H, W]
+                        elif task == 'depth':
+                            if pred_bidir_depth and refine_iter_idx == 0:
+                                intrinsics_curr = intrinsics_curr.repeat(2, 1, 1)
+                                pose = torch.cat((pose, torch.inverse(pose)), dim=0)
+
+                                feature0_ori, feature1_ori = torch.cat((feature0_ori, feature1_ori),
+                                                                       dim=0), torch.cat((feature1_ori,
+                                                                                          feature0_ori), dim=0)
+
+                            flow_from_depth = compute_flow_with_depth_pose(1. / flow.squeeze(1),
+                                                                           intrinsics_curr,
+                                                                           extrinsics_rel=pose,
+                                                                           )
+
+                            correlation = local_correlation_with_flow(
+                                feature0_ori,
+                                feature1_ori,
+                                flow=flow_from_depth,
+                                local_radius=4,
+                            )  # [B, (2R+1)^2, H, W]
+
+                        else:
+                            correlation = local_correlation_with_flow(
+                                feature0_ori,
+                                feature1_ori,
+                                flow=flow,
+                                local_radius=4,
+                            )  # [B, (2R+1)^2, H, W]
+
+                        proj = self.refine_proj(feature0)
+
+                        net, inp = torch.chunk(proj, chunks=2, dim=1)
+
+                        net = torch.tanh(net)
+                        inp = torch.relu(inp)
+
+                        net, up_mask, residual_flow = self.refine(net, inp, correlation, flow.clone(),
+                                                                  )
+
+                        if task == 'depth':
+                            flow = (flow - residual_flow).clamp(min=min_depth, max=max_depth)
+                        else:
+                            flow = flow + residual_flow
+
+                        if task == 'stereo':
+                            flow = flow.clamp(min=0)  # positive
+
+                        if self.training or refine_iter_idx == num_reg_refine - 1:
+                            if task == 'depth':
+                                if refine_iter_idx < num_reg_refine - 1:
+                                    # bilinear upsampling
+                                    flow_up = self.upsample_flow(flow, feature0, bilinear=True,
+                                                                 upsample_factor=upsample_factor,
+                                                                 is_depth=True)
+                                else:
+                                    # last one convex upsampling
+                                    # NOTE: clamp depth due to the zero padding in the unfold in the convex upsampling
+                                    # pad depth to 2 channels as flow
+                                    depth_pad = torch.cat((flow, torch.zeros_like(flow)), dim=1)  # [B, 2, H, W]
+                                    depth_up_pad = self.upsample_flow(depth_pad, feature0,
+                                                                      is_depth=True).clamp(min=min_depth,
+                                                                                           max=max_depth)
+                                    flow_up = depth_up_pad[:, :1]  # [B, 1, H, W]
+
+                            else:
+                                flow_up = upsample_flow_with_mask(flow, up_mask, upsample_factor=self.upsample_factor,
+                                                                  is_depth=task == 'depth')
+
+                            flow_preds.append(flow_up)
+
+        if task == 'stereo':
+            for i in range(len(flow_preds)):
+                flow_preds[i] = flow_preds[i].squeeze(1)  # [B, H, W]
+
+        # convert inverse depth to depth
+        if task == 'depth':
+            for i in range(len(flow_preds)):
+                flow_preds[i] = 1. / flow_preds[i].squeeze(1)  # [B, H, W]
+
+        results_dict.update({'flow_preds': flow_preds})
+
+        return results_dict

unimatch/utils.py

@@ -0,0 +1,216 @@
+import torch
+import torch.nn.functional as F
+from .position import PositionEmbeddingSine
+
+
+def generate_window_grid(h_min, h_max, w_min, w_max, len_h, len_w, device=None):
+    assert device is not None
+
+    x, y = torch.meshgrid([torch.linspace(w_min, w_max, len_w, device=device),
+                           torch.linspace(h_min, h_max, len_h, device=device)],
+                          )
+    grid = torch.stack((x, y), -1).transpose(0, 1).float()  # [H, W, 2]
+
+    return grid
+
+
+def normalize_coords(coords, h, w):
+    # coords: [B, H, W, 2]
+    c = torch.Tensor([(w - 1) / 2., (h - 1) / 2.]).float().to(coords.device)
+    return (coords - c) / c  # [-1, 1]
+
+
+def normalize_img(img0, img1):
+    # loaded images are in [0, 255]
+    # normalize by ImageNet mean and std
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(img1.device)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(img1.device)
+    img0 = (img0 / 255. - mean) / std
+    img1 = (img1 / 255. - mean) / std
+
+    return img0, img1
+
+
+def split_feature(feature,
+                  num_splits=2,
+                  channel_last=False,
+                  ):
+    if channel_last:  # [B, H, W, C]
+        b, h, w, c = feature.size()
+        assert h % num_splits == 0 and w % num_splits == 0
+
+        b_new = b * num_splits * num_splits
+        h_new = h // num_splits
+        w_new = w // num_splits
+
+        feature = feature.view(b, num_splits, h // num_splits, num_splits, w // num_splits, c
+                               ).permute(0, 1, 3, 2, 4, 5).reshape(b_new, h_new, w_new, c)  # [B*K*K, H/K, W/K, C]
+    else:  # [B, C, H, W]
+        b, c, h, w = feature.size()
+        assert h % num_splits == 0 and w % num_splits == 0
+
+        b_new = b * num_splits * num_splits
+        h_new = h // num_splits
+        w_new = w // num_splits
+
+        feature = feature.view(b, c, num_splits, h // num_splits, num_splits, w // num_splits
+                               ).permute(0, 2, 4, 1, 3, 5).reshape(b_new, c, h_new, w_new)  # [B*K*K, C, H/K, W/K]
+
+    return feature
+
+
+def merge_splits(splits,
+                 num_splits=2,
+                 channel_last=False,
+                 ):
+    if channel_last:  # [B*K*K, H/K, W/K, C]
+        b, h, w, c = splits.size()
+        new_b = b // num_splits // num_splits
+
+        splits = splits.view(new_b, num_splits, num_splits, h, w, c)
+        merge = splits.permute(0, 1, 3, 2, 4, 5).contiguous().view(
+            new_b, num_splits * h, num_splits * w, c)  # [B, H, W, C]
+    else:  # [B*K*K, C, H/K, W/K]
+        b, c, h, w = splits.size()
+        new_b = b // num_splits // num_splits
+
+        splits = splits.view(new_b, num_splits, num_splits, c, h, w)
+        merge = splits.permute(0, 3, 1, 4, 2, 5).contiguous().view(
+            new_b, c, num_splits * h, num_splits * w)  # [B, C, H, W]
+
+    return merge
+
+
+def generate_shift_window_attn_mask(input_resolution, window_size_h, window_size_w,
+                                    shift_size_h, shift_size_w, device=torch.device('cuda')):
+    # ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
+    # calculate attention mask for SW-MSA
+    h, w = input_resolution
+    img_mask = torch.zeros((1, h, w, 1)).to(device)  # 1 H W 1
+    h_slices = (slice(0, -window_size_h),
+                slice(-window_size_h, -shift_size_h),
+                slice(-shift_size_h, None))
+    w_slices = (slice(0, -window_size_w),
+                slice(-window_size_w, -shift_size_w),
+                slice(-shift_size_w, None))
+    cnt = 0
+    for h in h_slices:
+        for w in w_slices:
+            img_mask[:, h, w, :] = cnt
+            cnt += 1
+
+    mask_windows = split_feature(img_mask, num_splits=input_resolution[-1] // window_size_w, channel_last=True)
+
+    mask_windows = mask_windows.view(-1, window_size_h * window_size_w)
+    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+
+    return attn_mask
+
+
+def feature_add_position(feature0, feature1, attn_splits, feature_channels):
+    pos_enc = PositionEmbeddingSine(num_pos_feats=feature_channels // 2)
+
+    if attn_splits > 1:  # add position in splited window
+        feature0_splits = split_feature(feature0, num_splits=attn_splits)
+        feature1_splits = split_feature(feature1, num_splits=attn_splits)
+
+        position = pos_enc(feature0_splits)
+
+        feature0_splits = feature0_splits + position
+        feature1_splits = feature1_splits + position
+
+        feature0 = merge_splits(feature0_splits, num_splits=attn_splits)
+        feature1 = merge_splits(feature1_splits, num_splits=attn_splits)
+    else:
+        position = pos_enc(feature0)
+
+        feature0 = feature0 + position
+        feature1 = feature1 + position
+
+    return feature0, feature1
+
+
+def upsample_flow_with_mask(flow, up_mask, upsample_factor,
+                            is_depth=False):
+    # convex upsampling following raft
+
+    mask = up_mask
+    b, flow_channel, h, w = flow.shape
+    mask = mask.view(b, 1, 9, upsample_factor, upsample_factor, h, w)  # [B, 1, 9, K, K, H, W]
+    mask = torch.softmax(mask, dim=2)
+
+    multiplier = 1 if is_depth else upsample_factor
+    up_flow = F.unfold(multiplier * flow, [3, 3], padding=1)
+    up_flow = up_flow.view(b, flow_channel, 9, 1, 1, h, w)  # [B, 2, 9, 1, 1, H, W]
+
+    up_flow = torch.sum(mask * up_flow, dim=2)  # [B, 2, K, K, H, W]
+    up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)  # [B, 2, K, H, K, W]
+    up_flow = up_flow.reshape(b, flow_channel, upsample_factor * h,
+                              upsample_factor * w)  # [B, 2, K*H, K*W]
+
+    return up_flow
+
+
+def split_feature_1d(feature,
+                     num_splits=2,
+                     ):
+    # feature: [B, W, C]
+    b, w, c = feature.size()
+    assert w % num_splits == 0
+
+    b_new = b * num_splits
+    w_new = w // num_splits
+
+    feature = feature.view(b, num_splits, w // num_splits, c
+                           ).view(b_new, w_new, c)  # [B*K, W/K, C]
+
+    return feature
+
+
+def merge_splits_1d(splits,
+                    h,
+                    num_splits=2,
+                    ):
+    b, w, c = splits.size()
+    new_b = b // num_splits // h
+
+    splits = splits.view(new_b, h, num_splits, w, c)
+    merge = splits.view(
+        new_b, h, num_splits * w, c)  # [B, H, W, C]
+
+    return merge
+
+
+def window_partition_1d(x, window_size_w):
+    """
+    Args:
+        x: (B, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, C)
+    """
+    B, W, C = x.shape
+    x = x.view(B, W // window_size_w, window_size_w, C).view(-1, window_size_w, C)
+    return x
+
+
+def generate_shift_window_attn_mask_1d(input_w, window_size_w,
+                                       shift_size_w, device=torch.device('cuda')):
+    # calculate attention mask for SW-MSA
+    img_mask = torch.zeros((1, input_w, 1)).to(device)  # 1 W 1
+    w_slices = (slice(0, -window_size_w),
+                slice(-window_size_w, -shift_size_w),
+                slice(-shift_size_w, None))
+    cnt = 0
+    for w in w_slices:
+        img_mask[:, w, :] = cnt
+        cnt += 1
+
+    mask_windows = window_partition_1d(img_mask, window_size_w)  # nW, window_size, 1
+    mask_windows = mask_windows.view(-1, window_size_w)
+    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)  # nW, window_size, window_size
+    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+
+    return attn_mask

utils/flow_viz.py

@@ -0,0 +1,290 @@
+# MIT License
+#
+# Copyright (c) 2018 Tom Runia
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to conditions.
+#
+# Author: Tom Runia
+# Date Created: 2018-08-03
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from PIL import Image
+
+
+def make_colorwheel():
+    '''
+    Generates a color wheel for optical flow visualization as presented in:
+        Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
+        URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf
+    According to the C++ source code of Daniel Scharstein
+    According to the Matlab source code of Deqing Sun
+    '''
+
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros((ncols, 3))
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.floor(255 * np.arange(0, RY) / RY)
+    col = col + RY
+    # YG
+    colorwheel[col:col + YG, 0] = 255 - np.floor(255 * np.arange(0, YG) / YG)
+    colorwheel[col:col + YG, 1] = 255
+    col = col + YG
+    # GC
+    colorwheel[col:col + GC, 1] = 255
+    colorwheel[col:col + GC, 2] = np.floor(255 * np.arange(0, GC) / GC)
+    col = col + GC
+    # CB
+    colorwheel[col:col + CB, 1] = 255 - np.floor(255 * np.arange(CB) / CB)
+    colorwheel[col:col + CB, 2] = 255
+    col = col + CB
+    # BM
+    colorwheel[col:col + BM, 2] = 255
+    colorwheel[col:col + BM, 0] = np.floor(255 * np.arange(0, BM) / BM)
+    col = col + BM
+    # MR
+    colorwheel[col:col + MR, 2] = 255 - np.floor(255 * np.arange(MR) / MR)
+    colorwheel[col:col + MR, 0] = 255
+    return colorwheel
+
+
+def flow_compute_color(u, v, convert_to_bgr=False):
+    '''
+    Applies the flow color wheel to (possibly clipped) flow components u and v.
+    According to the C++ source code of Daniel Scharstein
+    According to the Matlab source code of Deqing Sun
+    :param u: np.ndarray, input horizontal flow
+    :param v: np.ndarray, input vertical flow
+    :param convert_to_bgr: bool, whether to change ordering and output BGR instead of RGB
+    :return:
+    '''
+
+    flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
+
+    colorwheel = make_colorwheel()  # shape [55x3]
+    ncols = colorwheel.shape[0]
+
+    rad = np.sqrt(np.square(u) + np.square(v))
+    a = np.arctan2(-v, -u) / np.pi
+
+    fk = (a + 1) / 2 * (ncols - 1) + 1
+    k0 = np.floor(fk).astype(np.int32)
+    k1 = k0 + 1
+    k1[k1 == ncols] = 1
+    f = fk - k0
+
+    for i in range(colorwheel.shape[1]):
+        tmp = colorwheel[:, i]
+        col0 = tmp[k0] / 255.0
+        col1 = tmp[k1] / 255.0
+        col = (1 - f) * col0 + f * col1
+
+        idx = (rad <= 1)
+        col[idx] = 1 - rad[idx] * (1 - col[idx])
+        col[~idx] = col[~idx] * 0.75  # out of range?
+
+        # Note the 2-i => BGR instead of RGB
+        ch_idx = 2 - i if convert_to_bgr else i
+        flow_image[:, :, ch_idx] = np.floor(255 * col)
+
+    return flow_image
+
+
+def flow_to_color(flow_uv, clip_flow=None, convert_to_bgr=False):
+    '''
+    Expects a two dimensional flow image of shape [H,W,2]
+    According to the C++ source code of Daniel Scharstein
+    According to the Matlab source code of Deqing Sun
+    :param flow_uv: np.ndarray of shape [H,W,2]
+    :param clip_flow: float, maximum clipping value for flow
+    :return:
+    '''
+
+    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
+    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
+
+    if clip_flow is not None:
+        flow_uv = np.clip(flow_uv, 0, clip_flow)
+
+    u = flow_uv[:, :, 0]
+    v = flow_uv[:, :, 1]
+
+    rad = np.sqrt(np.square(u) + np.square(v))
+    rad_max = np.max(rad)
+
+    epsilon = 1e-5
+    u = u / (rad_max + epsilon)
+    v = v / (rad_max + epsilon)
+
+    return flow_compute_color(u, v, convert_to_bgr)
+
+
+UNKNOWN_FLOW_THRESH = 1e7
+SMALLFLOW = 0.0
+LARGEFLOW = 1e8
+
+
+def make_color_wheel():
+    """
+    Generate color wheel according Middlebury color code
+    :return: Color wheel
+    """
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+
+    colorwheel = np.zeros([ncols, 3])
+
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY))
+    col += RY
+
+    # YG
+    colorwheel[col:col + YG, 0] = 255 - np.transpose(np.floor(255 * np.arange(0, YG) / YG))
+    colorwheel[col:col + YG, 1] = 255
+    col += YG
+
+    # GC
+    colorwheel[col:col + GC, 1] = 255
+    colorwheel[col:col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC))
+    col += GC
+
+    # CB
+    colorwheel[col:col + CB, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, CB) / CB))
+    colorwheel[col:col + CB, 2] = 255
+    col += CB
+
+    # BM
+    colorwheel[col:col + BM, 2] = 255
+    colorwheel[col:col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM))
+    col += + BM
+
+    # MR
+    colorwheel[col:col + MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
+    colorwheel[col:col + MR, 0] = 255
+
+    return colorwheel
+
+
+def compute_color(u, v):
+    """
+    compute optical flow color map
+    :param u: optical flow horizontal map
+    :param v: optical flow vertical map
+    :return: optical flow in color code
+    """
+    [h, w] = u.shape
+    img = np.zeros([h, w, 3])
+    nanIdx = np.isnan(u) | np.isnan(v)
+    u[nanIdx] = 0
+    v[nanIdx] = 0
+
+    colorwheel = make_color_wheel()
+    ncols = np.size(colorwheel, 0)
+
+    rad = np.sqrt(u ** 2 + v ** 2)
+
+    a = np.arctan2(-v, -u) / np.pi
+
+    fk = (a + 1) / 2 * (ncols - 1) + 1
+
+    k0 = np.floor(fk).astype(int)
+
+    k1 = k0 + 1
+    k1[k1 == ncols + 1] = 1
+    f = fk - k0
+
+    for i in range(0, np.size(colorwheel, 1)):
+        tmp = colorwheel[:, i]
+        col0 = tmp[k0 - 1] / 255
+        col1 = tmp[k1 - 1] / 255
+        col = (1 - f) * col0 + f * col1
+
+        idx = rad <= 1
+        col[idx] = 1 - rad[idx] * (1 - col[idx])
+        notidx = np.logical_not(idx)
+
+        col[notidx] *= 0.75
+        img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx)))
+
+    return img
+
+
+# from https://github.com/gengshan-y/VCN
+def flow_to_image(flow):
+    """
+    Convert flow into middlebury color code image
+    :param flow: optical flow map
+    :return: optical flow image in middlebury color
+    """
+    u = flow[:, :, 0]
+    v = flow[:, :, 1]
+
+    maxu = -999.
+    maxv = -999.
+    minu = 999.
+    minv = 999.
+
+    idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
+    u[idxUnknow] = 0
+    v[idxUnknow] = 0
+
+    maxu = max(maxu, np.max(u))
+    minu = min(minu, np.min(u))
+
+    maxv = max(maxv, np.max(v))
+    minv = min(minv, np.min(v))
+
+    rad = np.sqrt(u ** 2 + v ** 2)
+    maxrad = max(-1, np.max(rad))
+
+    u = u / (maxrad + np.finfo(float).eps)
+    v = v / (maxrad + np.finfo(float).eps)
+
+    img = compute_color(u, v)
+
+    idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
+    img[idx] = 0
+
+    return np.uint8(img)
+
+
+def save_vis_flow_tofile(flow, output_path):
+    vis_flow = flow_to_image(flow)
+    Image.fromarray(vis_flow).save(output_path)
+
+
+def flow_tensor_to_image(flow):
+    """Used for tensorboard visualization"""
+    flow = flow.permute(1, 2, 0)  # [H, W, 2]
+    flow = flow.detach().cpu().numpy()
+    flow = flow_to_image(flow)  # [H, W, 3]
+    flow = np.transpose(flow, (2, 0, 1))  # [3, H, W]
+
+    return flow

utils/visualization.py

@@ -0,0 +1,110 @@
+import torch
+import torch.utils.data
+import numpy as np
+import torchvision.utils as vutils
+import cv2
+from matplotlib.cm import get_cmap
+import matplotlib as mpl
+import matplotlib.cm as cm
+
+
+def vis_disparity(disp, return_rgb=False):
+    disp_vis = (disp - disp.min()) / (disp.max() - disp.min()) * 255.0
+    disp_vis = disp_vis.astype("uint8")
+    disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
+
+    if return_rgb:
+        disp_vis = cv2.cvtColor(disp_vis, cv2.COLOR_BGR2RGB)
+
+    return disp_vis
+
+
+def gen_error_colormap():
+    cols = np.array(
+        [[0 / 3.0, 0.1875 / 3.0, 49, 54, 149],
+         [0.1875 / 3.0, 0.375 / 3.0, 69, 117, 180],
+         [0.375 / 3.0, 0.75 / 3.0, 116, 173, 209],
+         [0.75 / 3.0, 1.5 / 3.0, 171, 217, 233],
+         [1.5 / 3.0, 3 / 3.0, 224, 243, 248],
+         [3 / 3.0, 6 / 3.0, 254, 224, 144],
+         [6 / 3.0, 12 / 3.0, 253, 174, 97],
+         [12 / 3.0, 24 / 3.0, 244, 109, 67],
+         [24 / 3.0, 48 / 3.0, 215, 48, 39],
+         [48 / 3.0, np.inf, 165, 0, 38]], dtype=np.float32)
+    cols[:, 2: 5] /= 255.
+    return cols
+
+
+def disp_error_img(D_est_tensor, D_gt_tensor, abs_thres=3., rel_thres=0.05, dilate_radius=1):
+    D_gt_np = D_gt_tensor.detach().cpu().numpy()
+    D_est_np = D_est_tensor.detach().cpu().numpy()
+    B, H, W = D_gt_np.shape
+    # valid mask
+    mask = D_gt_np > 0
+    # error in percentage. When error <= 1, the pixel is valid since <= 3px & 5%
+    error = np.abs(D_gt_np - D_est_np)
+    error[np.logical_not(mask)] = 0
+    error[mask] = np.minimum(error[mask] / abs_thres, (error[mask] / D_gt_np[mask]) / rel_thres)
+    # get colormap
+    cols = gen_error_colormap()
+    # create error image
+    error_image = np.zeros([B, H, W, 3], dtype=np.float32)
+    for i in range(cols.shape[0]):
+        error_image[np.logical_and(error >= cols[i][0], error < cols[i][1])] = cols[i, 2:]
+    # TODO: imdilate
+    # error_image = cv2.imdilate(D_err, strel('disk', dilate_radius));
+    error_image[np.logical_not(mask)] = 0.
+    # show color tag in the top-left cornor of the image
+    for i in range(cols.shape[0]):
+        distance = 20
+        error_image[:, :10, i * distance:(i + 1) * distance, :] = cols[i, 2:]
+
+    return torch.from_numpy(np.ascontiguousarray(error_image.transpose([0, 3, 1, 2])))
+
+
+def save_images(logger, mode_tag, images_dict, global_step):
+    images_dict = tensor2numpy(images_dict)
+    for tag, values in images_dict.items():
+        if not isinstance(values, list) and not isinstance(values, tuple):
+            values = [values]
+        for idx, value in enumerate(values):
+            if len(value.shape) == 3:
+                value = value[:, np.newaxis, :, :]
+            value = value[:1]
+            value = torch.from_numpy(value)
+
+            image_name = '{}/{}'.format(mode_tag, tag)
+            if len(values) > 1:
+                image_name = image_name + "_" + str(idx)
+            logger.add_image(image_name, vutils.make_grid(value, padding=0, nrow=1, normalize=True, scale_each=True),
+                             global_step)
+
+
+def tensor2numpy(var_dict):
+    for key, vars in var_dict.items():
+        if isinstance(vars, np.ndarray):
+            var_dict[key] = vars
+        elif isinstance(vars, torch.Tensor):
+            var_dict[key] = vars.data.cpu().numpy()
+        else:
+            raise NotImplementedError("invalid input type for tensor2numpy")
+
+    return var_dict
+
+
+def viz_depth_tensor_from_monodepth2(disp, return_numpy=False, colormap='plasma'):
+    # visualize inverse depth
+    assert isinstance(disp, torch.Tensor)
+
+    disp = disp.numpy()
+    vmax = np.percentile(disp, 95)
+    normalizer = mpl.colors.Normalize(vmin=disp.min(), vmax=vmax)
+    mapper = cm.ScalarMappable(norm=normalizer, cmap=colormap)
+    colormapped_im = (mapper.to_rgba(disp)[:, :, :3] * 255).astype(np.uint8)  # [H, W, 3]
+
+    if return_numpy:
+        return colormapped_im
+
+    viz = torch.from_numpy(colormapped_im).permute(2, 0, 1)  # [3, H, W]
+
+    return viz