docstring for WarpLayer
| 6 | from torch.autograd import Variable |
| 7 | |
| 8 | class ForwardWarp(nn.Module): |
| 9 | """docstring for WarpLayer""" |
| 10 | def __init__(self,): |
| 11 | super(ForwardWarp, self).__init__() |
| 12 | |
| 13 | def forward(self, img, flo): |
| 14 | """ |
| 15 | -img: image (N, C, H, W) |
| 16 | -flo: optical flow (N, 2, H, W) |
| 17 | elements of flo is in [0, H] and [0, W] for dx, dy |
| 18 | |
| 19 | """ |
| 20 | |
| 21 | |
| 22 | # (x1, y1) (x1, y2) |
| 23 | # +---------------+ |
| 24 | # | | |
| 25 | # | o(x, y) | |
| 26 | # | | |
| 27 | # | | |
| 28 | # | | |
| 29 | # | | |
| 30 | # +---------------+ |
| 31 | # (x2, y1) (x2, y2) |
| 32 | |
| 33 | |
| 34 | N, C, _, _ = img.size() |
| 35 | |
| 36 | # translate start-point optical flow to end-point optical flow |
| 37 | y = flo[:, 0:1 :, :] |
| 38 | x = flo[:, 1:2, :, :] |
| 39 | |
| 40 | x = x.repeat(1, C, 1, 1) |
| 41 | y = y.repeat(1, C, 1, 1) |
| 42 | |
| 43 | # Four point of square (x1, y1), (x1, y2), (x2, y1), (y2, y2) |
| 44 | x1 = torch.floor(x) |
| 45 | x2 = x1 + 1 |
| 46 | y1 = torch.floor(y) |
| 47 | y2 = y1 + 1 |
| 48 | |
| 49 | # firstly, get gaussian weights |
| 50 | w11, w12, w21, w22 = self.get_gaussian_weights(x, y, x1, x2, y1, y2) |
| 51 | |
| 52 | # secondly, sample each weighted corner |
| 53 | img11, o11 = self.sample_one(img, x1, y1, w11) |
| 54 | img12, o12 = self.sample_one(img, x1, y2, w12) |
| 55 | img21, o21 = self.sample_one(img, x2, y1, w21) |
| 56 | img22, o22 = self.sample_one(img, x2, y2, w22) |
| 57 | |
| 58 | |
| 59 | imgw = img11 + img12 + img21 + img22 |
| 60 | o = o11 + o12 + o21 + o22 |
| 61 | |
| 62 | return imgw, o |
| 63 | |
| 64 | |
| 65 | def get_gaussian_weights(self, x, y, x1, x2, y1, y2): |
no outgoing calls
no test coverage detected