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