57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from utils.utils import bilinear_sampler, coords_grid
|
|
|
|
|
|
class CorrBlock:
|
|
def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
|
|
self.num_levels = num_levels
|
|
self.radius = radius
|
|
self.corr_pyramid = []
|
|
|
|
# all pairs correlation
|
|
corr = CorrBlock.corr(fmap1, fmap2)
|
|
|
|
batch, h1, w1, dim, h2, w2 = corr.shape
|
|
corr = corr.reshape(batch*h1*w1, dim, h2, w2)
|
|
|
|
self.corr_pyramid.append(corr)
|
|
for i in range(self.num_levels-1):
|
|
corr = F.avg_pool2d(corr, 2, stride=2)
|
|
self.corr_pyramid.append(corr)
|
|
|
|
def __call__(self, coords):
|
|
r = self.radius
|
|
coords = coords.permute(0, 2, 3, 1)
|
|
batch, h1, w1, _ = coords.shape
|
|
|
|
out_pyramid = []
|
|
for i in range(self.num_levels):
|
|
corr = self.corr_pyramid[i]
|
|
dx = torch.linspace(-r, r, 2*r+1)
|
|
dy = torch.linspace(-r, r, 2*r+1)
|
|
delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device)
|
|
|
|
centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i
|
|
delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
|
|
coords_lvl = centroid_lvl + delta_lvl
|
|
|
|
corr = bilinear_sampler(corr, coords_lvl)
|
|
corr = corr.view(batch, h1, w1, -1)
|
|
out_pyramid.append(corr)
|
|
|
|
out = torch.cat(out_pyramid, dim=-1)
|
|
return out.permute(0, 3, 1, 2).contiguous().float()
|
|
|
|
|
|
@staticmethod
|
|
def corr(fmap1, fmap2):
|
|
batch, dim, ht, wd = fmap1.shape
|
|
fmap1 = fmap1.view(batch, dim, ht*wd)
|
|
fmap2 = fmap2.view(batch, dim, ht*wd)
|
|
|
|
corr = torch.matmul(fmap1.transpose(1,2), fmap2)
|
|
corr = corr.view(batch, ht, wd, 1, ht, wd)
|
|
return corr / torch.sqrt(torch.tensor(dim).float())
|
|
|