| 521 | |
| 522 | |
| 523 | class TpsGridGen(Module): |
| 524 | def __init__(self, out_h=240, out_w=240, use_regular_grid=True, grid_size=3, reg_factor=0, use_cuda=True): |
| 525 | super(TpsGridGen, self).__init__() |
| 526 | self.out_h, self.out_w = out_h, out_w |
| 527 | self.reg_factor = reg_factor |
| 528 | self.use_cuda = use_cuda |
| 529 | |
| 530 | # create grid in numpy |
| 531 | # self.grid = np.zeros( [self.out_h, self.out_w, 3], dtype=np.float32) |
| 532 | # sampling grid with dim-0 coords (Y) |
| 533 | self.grid_X, self.grid_Y = np.meshgrid(np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h)) |
| 534 | # grid_X,grid_Y: size [1,H,W,1,1] |
| 535 | self.grid_X = torch.FloatTensor(self.grid_X).unsqueeze(0).unsqueeze(3) |
| 536 | self.grid_Y = torch.FloatTensor(self.grid_Y).unsqueeze(0).unsqueeze(3) |
| 537 | self.grid_X = Variable(self.grid_X, requires_grad=False) |
| 538 | self.grid_Y = Variable(self.grid_Y, requires_grad=False) |
| 539 | if use_cuda: |
| 540 | self.grid_X = self.grid_X.cuda() |
| 541 | self.grid_Y = self.grid_Y.cuda() |
| 542 | |
| 543 | # initialize regular grid for control points P_i |
| 544 | if use_regular_grid: |
| 545 | axis_coords = np.linspace(-1, 1, grid_size) |
| 546 | self.N = grid_size * grid_size |
| 547 | P_Y, P_X = np.meshgrid(axis_coords, axis_coords) |
| 548 | P_X = np.reshape(P_X, (-1, 1)) # size (N,1) |
| 549 | P_Y = np.reshape(P_Y, (-1, 1)) # size (N,1) |
| 550 | P_X = torch.FloatTensor(P_X) |
| 551 | P_Y = torch.FloatTensor(P_Y) |
| 552 | self.Li = Variable(self.compute_L_inverse(P_X, P_Y).unsqueeze(0), requires_grad=False) |
| 553 | self.P_X = P_X.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(0, 4) |
| 554 | self.P_Y = P_Y.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(0, 4) |
| 555 | self.P_X = Variable(self.P_X, requires_grad=False) |
| 556 | self.P_Y = Variable(self.P_Y, requires_grad=False) |
| 557 | if use_cuda: |
| 558 | self.P_X = self.P_X.cuda() |
| 559 | self.P_Y = self.P_Y.cuda() |
| 560 | |
| 561 | def forward(self, theta): |
| 562 | warped_grid = self.apply_transformation(theta, torch.cat((self.grid_X, self.grid_Y), 3)) |
| 563 | |
| 564 | return warped_grid |
| 565 | |
| 566 | def compute_L_inverse(self, X, Y): |
| 567 | N = X.size()[0] # num of points (along dim 0) |
| 568 | # construct matrix K |
| 569 | Xmat = X.expand(N, N) |
| 570 | Ymat = Y.expand(N, N) |
| 571 | P_dist_squared = torch.pow(Xmat - Xmat.transpose(0, 1), 2) + torch.pow(Ymat - Ymat.transpose(0, 1), 2) |
| 572 | P_dist_squared[P_dist_squared == 0] = 1 # make diagonal 1 to avoid NaN in log computation |
| 573 | K = torch.mul(P_dist_squared, torch.log(P_dist_squared)) |
| 574 | if self.reg_factor != 0: |
| 575 | K += torch.eye(K.size(0), K.size(1)) * self.reg_factor |
| 576 | # construct matrix L |
| 577 | O = torch.FloatTensor(N, 1).fill_(1) |
| 578 | Z = torch.FloatTensor(3, 3).fill_(0) |
| 579 | P = torch.cat((O, X, Y), 1) |
| 580 | L = torch.cat((torch.cat((K, P), 1), torch.cat((P.transpose(0, 1), Z), 1)), 0) |