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