| 400 | |
| 401 | |
| 402 | class AffineGridGenV2(Module): |
| 403 | def __init__(self, out_h=240, out_w=240, use_cuda=True): |
| 404 | super(AffineGridGenV2, self).__init__() |
| 405 | self.out_h, self.out_w = out_h, out_w |
| 406 | self.use_cuda = use_cuda |
| 407 | |
| 408 | # create grid in numpy |
| 409 | # self.grid = np.zeros( [self.out_h, self.out_w, 3], dtype=np.float32) |
| 410 | # sampling grid with dim-0 coords (Y) |
| 411 | self.grid_X, self.grid_Y = np.meshgrid(np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h)) |
| 412 | # grid_X,grid_Y: size [1,H,W,1,1] |
| 413 | self.grid_X = torch.FloatTensor(self.grid_X).unsqueeze(0).unsqueeze(3) |
| 414 | self.grid_Y = torch.FloatTensor(self.grid_Y).unsqueeze(0).unsqueeze(3) |
| 415 | self.grid_X = Variable(self.grid_X, requires_grad=False) |
| 416 | self.grid_Y = Variable(self.grid_Y, requires_grad=False) |
| 417 | if use_cuda: |
| 418 | self.grid_X = self.grid_X.cuda() |
| 419 | self.grid_Y = self.grid_Y.cuda() |
| 420 | |
| 421 | def forward(self, theta): |
| 422 | b = theta.size(0) |
| 423 | if not theta.size() == (b, 6): |
| 424 | theta = theta.view(b, 6) |
| 425 | theta = theta.contiguous() |
| 426 | |
| 427 | t0 = theta[:, 0].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 428 | t1 = theta[:, 1].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 429 | t2 = theta[:, 2].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 430 | t3 = theta[:, 3].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 431 | t4 = theta[:, 4].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 432 | t5 = theta[:, 5].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 433 | |
| 434 | grid_X = expand_dim(self.grid_X, 0, b) |
| 435 | grid_Y = expand_dim(self.grid_Y, 0, b) |
| 436 | grid_Xp = grid_X * t0 + grid_Y * t1 + t2 |
| 437 | grid_Yp = grid_X * t3 + grid_Y * t4 + t5 |
| 438 | |
| 439 | return torch.cat((grid_Xp, grid_Yp), 3) |
| 440 | |
| 441 | |
| 442 | class HomographyGridGen(Module): |