(self, image_batch, theta_aff, theta_aff_tps, use_cuda=True)
| 34 | use_cuda=use_cuda) |
| 35 | |
| 36 | def __call__(self, image_batch, theta_aff, theta_aff_tps, use_cuda=True): |
| 37 | sampling_grid_aff = self.affTnf(image_batch=None, |
| 38 | theta_batch=theta_aff.view(-1, 2, 3), |
| 39 | return_sampling_grid=True, |
| 40 | return_warped_image=False) |
| 41 | |
| 42 | sampling_grid_aff_tps = self.tpsTnf(image_batch=None, |
| 43 | theta_batch=theta_aff_tps, |
| 44 | return_sampling_grid=True, |
| 45 | return_warped_image=False) |
| 46 | |
| 47 | if self.padding_crop_factor is not None: |
| 48 | sampling_grid_aff_tps = sampling_grid_aff_tps * self.padding_crop_factor |
| 49 | |
| 50 | # put 1e10 value in region out of bounds of sampling_grid_aff |
| 51 | in_bound_mask_aff = ((sampling_grid_aff[:, :, :, 0] > -1) * (sampling_grid_aff[:, :, :, 0] < 1) * ( |
| 52 | sampling_grid_aff[:, :, :, 1] > -1) * (sampling_grid_aff[:, :, :, 1] < 1)).unsqueeze(3) |
| 53 | in_bound_mask_aff = in_bound_mask_aff.expand_as(sampling_grid_aff) |
| 54 | sampling_grid_aff = torch.mul(in_bound_mask_aff.float(), sampling_grid_aff) |
| 55 | sampling_grid_aff = torch.add((in_bound_mask_aff.float() - 1) * (1e10), sampling_grid_aff) |
| 56 | |
| 57 | # compose transformations |
| 58 | sampling_grid_aff_tps_comp = F.grid_sample(sampling_grid_aff.transpose(2, 3).transpose(1, 2), |
| 59 | sampling_grid_aff_tps, align_corners=True).transpose(1, 2).transpose(2, 3) |
| 60 | |
| 61 | # put 1e10 value in region out of bounds of sampling_grid_aff_tps_comp |
| 62 | in_bound_mask_aff_tps = ((sampling_grid_aff_tps[:, :, :, 0] > -1) * (sampling_grid_aff_tps[:, :, :, 0] < 1) * ( |
| 63 | sampling_grid_aff_tps[:, :, :, 1] > -1) * (sampling_grid_aff_tps[:, :, :, 1] < 1)).unsqueeze(3) |
| 64 | in_bound_mask_aff_tps = in_bound_mask_aff_tps.expand_as(sampling_grid_aff_tps_comp) |
| 65 | sampling_grid_aff_tps_comp = torch.mul(in_bound_mask_aff_tps.float(), sampling_grid_aff_tps_comp) |
| 66 | sampling_grid_aff_tps_comp = torch.add((in_bound_mask_aff_tps.float() - 1) * (1e10), sampling_grid_aff_tps_comp) |
| 67 | |
| 68 | # sample transformed image |
| 69 | warped_image_batch = F.grid_sample(image_batch, sampling_grid_aff_tps_comp, align_corners=True) |
| 70 | |
| 71 | return warped_image_batch |
| 72 | |
| 73 | |
| 74 | class GeometricTnf(object): |
nothing calls this directly
no outgoing calls
no test coverage detected