(self, out_h=240, out_w=240, use_regular_grid=True, grid_size=3, reg_factor=0, use_cuda=True)
| 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)) |
nothing calls this directly
no test coverage detected