Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination
(self, flow, mask)
| 75 | return coords0, coords1 |
| 76 | |
| 77 | def upsample_flow(self, flow, mask): |
| 78 | """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """ |
| 79 | N, _, H, W = flow.shape |
| 80 | mask = mask.view(N, 1, 9, 8, 8, H, W) |
| 81 | mask = torch.softmax(mask, dim=2) |
| 82 | |
| 83 | up_flow = F.unfold(8 * flow, [3,3], padding=1) |
| 84 | up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) |
| 85 | |
| 86 | up_flow = torch.sum(mask * up_flow, dim=2) |
| 87 | up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) |
| 88 | return up_flow.reshape(N, 2, 8*H, 8*W) |
| 89 | |
| 90 | |
| 91 | def forward(self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False): |