| 440 | |
| 441 | |
| 442 | class HomographyGridGen(Module): |
| 443 | def __init__(self, out_h=240, out_w=240, use_cuda=True): |
| 444 | super(HomographyGridGen, self).__init__() |
| 445 | self.out_h, self.out_w = out_h, out_w |
| 446 | self.use_cuda = use_cuda |
| 447 | |
| 448 | # create grid in numpy |
| 449 | # self.grid = np.zeros( [self.out_h, self.out_w, 3], dtype=np.float32) |
| 450 | # sampling grid with dim-0 coords (Y) |
| 451 | self.grid_X, self.grid_Y = np.meshgrid(np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h)) |
| 452 | # grid_X,grid_Y: size [1,H,W,1,1] |
| 453 | self.grid_X = torch.FloatTensor(self.grid_X).unsqueeze(0).unsqueeze(3) |
| 454 | self.grid_Y = torch.FloatTensor(self.grid_Y).unsqueeze(0).unsqueeze(3) |
| 455 | self.grid_X = Variable(self.grid_X, requires_grad=False) |
| 456 | self.grid_Y = Variable(self.grid_Y, requires_grad=False) |
| 457 | if use_cuda: |
| 458 | self.grid_X = self.grid_X.cuda() |
| 459 | self.grid_Y = self.grid_Y.cuda() |
| 460 | |
| 461 | def forward(self, theta): |
| 462 | b = theta.size(0) |
| 463 | if theta.size(1) == 9: |
| 464 | H = theta |
| 465 | else: |
| 466 | H = homography_mat_from_4_pts(theta) |
| 467 | h0 = H[:, 0].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 468 | h1 = H[:, 1].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 469 | h2 = H[:, 2].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 470 | h3 = H[:, 3].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 471 | h4 = H[:, 4].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 472 | h5 = H[:, 5].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 473 | h6 = H[:, 6].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 474 | h7 = H[:, 7].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 475 | h8 = H[:, 8].unsqueeze(1).unsqueeze(2).unsqueeze(3) |
| 476 | |
| 477 | grid_X = expand_dim(self.grid_X, 0, b) |
| 478 | grid_Y = expand_dim(self.grid_Y, 0, b) |
| 479 | |
| 480 | grid_Xp = grid_X * h0 + grid_Y * h1 + h2 |
| 481 | grid_Yp = grid_X * h3 + grid_Y * h4 + h5 |
| 482 | k = grid_X * h6 + grid_Y * h7 + h8 |
| 483 | |
| 484 | grid_Xp /= k |
| 485 | grid_Yp /= k |
| 486 | |
| 487 | return torch.cat((grid_Xp, grid_Yp), 3) |
| 488 | |
| 489 | |
| 490 | def homography_mat_from_4_pts(theta): |