(self, coords)
| 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): |
nothing calls this directly
no test coverage detected