def forward(self, x):
if self.fine_grid is None:
# not simplified (training/validation)
# generate coarse affine and TPS grids
coarse_affine_grid = F.affine_grid(
self.affine_mat,
torch.Size([1, x.shape[1], x.shape[2], x.shape[3]])).permute(
(0, 3, 1, 2))
coarse_tps_grid = pytorch_tps.tps_grid(
self.theta, self.ctrl_pts, (1, x.size()[1]) + self.out_size)
# use TPS grid to sample affine grid
tps_grid = F.grid_sample(coarse_affine_grid,
coarse_tps_grid).repeat(
x.shape[0], 1, 1, 1)
# refine TPS grid using grid refinement net and save it to self.fine_grid
if self.with_refine:
fine_grid = torch.clamp(self.grid_refine_net(tps_grid) +
tps_grid,
min=-1,
max=1).permute((0, 2, 3, 1))
else:
fine_grid = torch.clamp(tps_grid, min=-1, max=1).permute(
(0, 2, 3, 1))
else:
# simplified (testing)
fine_grid = self.fine_grid.repeat(x.shape[0], 1, 1, 1)
# warp
x = F.grid_sample(x, fine_grid)
return x
def simplify(self, x):
# generate coarse affine and TPS grids
coarse_affine_grid = F.affine_grid(
self.affine_mat,
torch.Size([1, x.shape[1], x.shape[2], x.shape[3]])).permute(
(0, 3, 1, 2))
coarse_tps_grid = pytorch_tps.tps_grid(
self.theta, self.ctrl_pts, (1, x.size()[1]) + self.out_size)
# use TPS grid to sample affine grid
tps_grid = F.grid_sample(coarse_affine_grid, coarse_tps_grid)
# refine TPS grid using grid refinement net and save it to self.fine_grid
if self.with_refine:
self.fine_grid = torch.clamp(self.grid_refine_net(tps_grid) +
tps_grid,
min=-1,
max=1).permute((0, 2, 3, 1))
else:
self.fine_grid = torch.clamp(tps_grid, min=-1, max=1).permute(
(0, 2, 3, 1))