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]
+