warp an image/tensor (im2) back to im1, according to the optical flow x: [B, C, H, W] (im2) flo: [B, 2, H, W] flow
(self, x, flo)
| 157 | |
| 158 | |
| 159 | def warp(self, x, flo): |
| 160 | """ |
| 161 | warp an image/tensor (im2) back to im1, according to the optical flow |
| 162 | |
| 163 | x: [B, C, H, W] (im2) |
| 164 | flo: [B, 2, H, W] flow |
| 165 | |
| 166 | """ |
| 167 | B, C, H, W = x.size() |
| 168 | # mesh grid |
| 169 | # xx = torch.arange(0, W).view(1,-1).cuda().repeat(H,1) |
| 170 | # yy = torch.arange(0, H).view(-1,1).cuda().repeat(1,W) |
| 171 | # xx = xx.view(1,1,H,W).repeat(B,1,1,1) |
| 172 | # yy = yy.view(1,1,H,W).repeat(B,1,1,1) |
| 173 | # grid = torch.cat((xx,yy),1).float() |
| 174 | |
| 175 | # # if x.is_cuda: |
| 176 | # # grid = grid.cuda() |
| 177 | # vgrid = Variable(grid) + flo |
| 178 | assert(B <= self.B_MAX and H <= self.H_MAX and W <= self.W_MAX) |
| 179 | vgrid = self.grid[:B,:,:H,:W] +flo |
| 180 | |
| 181 | # scale grid to [-1,1] |
| 182 | vgrid[:,0,:,:] = 2.0*vgrid[:,0,:,:].clone()/max(W-1,1)-1.0 |
| 183 | vgrid[:,1,:,:] = 2.0*vgrid[:,1,:,:].clone()/max(H-1,1)-1.0 |
| 184 | |
| 185 | |
| 186 | vgrid = vgrid.permute(0,2,3,1) |
| 187 | output = nn.functional.grid_sample(x, vgrid) |
| 188 | # mask = torch.autograd.Variable(torch.ones(x.size())).cuda() |
| 189 | mask = torch.autograd.Variable(torch.cuda.FloatTensor().resize_(x.size()).zero_() + 1, requires_grad = False) |
| 190 | mask = nn.functional.grid_sample(mask, vgrid) |
| 191 | |
| 192 | # if W==128: |
| 193 | # np.save('mask.npy', mask.cpu().data.numpy()) |
| 194 | # np.save('warp.npy', output.cpu().data.numpy()) |
| 195 | |
| 196 | mask[mask<0.9999] = 0 |
| 197 | mask[mask>0] = 1 |
| 198 | |
| 199 | return output*mask |
| 200 | |
| 201 | |
| 202 | def forward(self,x, output_more = False): |