| 10 | |
| 11 | |
| 12 | class CorrBlock: |
| 13 | def __init__(self, fmap1, fmap2, num_levels=4, radius=4): |
| 14 | self.num_levels = num_levels |
| 15 | self.radius = radius |
| 16 | self.corr_pyramid = [] |
| 17 | |
| 18 | # all pairs correlation |
| 19 | corr = CorrBlock.corr(fmap1, fmap2) |
| 20 | |
| 21 | batch, h1, w1, dim, h2, w2 = corr.shape |
| 22 | corr = corr.reshape(batch*h1*w1, dim, h2, w2) |
| 23 | |
| 24 | self.corr_pyramid.append(corr) |
| 25 | for i in range(self.num_levels-1): |
| 26 | corr = F.avg_pool2d(corr, 2, stride=2) |
| 27 | self.corr_pyramid.append(corr) |
| 28 | |
| 29 | def __call__(self, coords): |
| 30 | r = self.radius |
| 31 | coords = coords.permute(0, 2, 3, 1) |
| 32 | batch, h1, w1, _ = coords.shape |
| 33 | |
| 34 | out_pyramid = [] |
| 35 | for i in range(self.num_levels): |
| 36 | corr = self.corr_pyramid[i] |
| 37 | dx = torch.linspace(-r, r, 2*r+1) |
| 38 | dy = torch.linspace(-r, r, 2*r+1) |
| 39 | delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device) |
| 40 | |
| 41 | centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i |
| 42 | delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2) |
| 43 | coords_lvl = centroid_lvl + delta_lvl |
| 44 | corr = bilinear_sampler(corr, coords_lvl) |
| 45 | corr = corr.view(batch, h1, w1, -1) |
| 46 | out_pyramid.append(corr) |
| 47 | |
| 48 | out = torch.cat(out_pyramid, dim=-1) |
| 49 | return out.permute(0, 3, 1, 2).contiguous().float() |
| 50 | |
| 51 | @staticmethod |
| 52 | def corr(fmap1, fmap2): |
| 53 | batch, dim, ht, wd = fmap1.shape |
| 54 | fmap1 = fmap1.view(batch, dim, ht*wd) |
| 55 | fmap2 = fmap2.view(batch, dim, ht*wd) |
| 56 | |
| 57 | corr = torch.matmul(fmap1.transpose(1,2), fmap2) |
| 58 | corr = corr.view(batch, ht, wd, 1, ht, wd) |
| 59 | return corr / torch.sqrt(torch.tensor(dim).float()) |
| 60 | |
| 61 | |
| 62 | class AlternateCorrBlock: |