Input: im: Torch.tensor[C, H, W] tnf_type: 'random', 'affine', 'hom', 'tps' Output: im: Torch.tensor[C, H, W] grid: Torch.tensor[H, W, 2]
(self, im, tnf_type, theta=None)
| 13 | self.theta_generator = {tnf: Theta_gen(geometric_model=tnf, output_size=(H, W)) for tnf in self.tnf_list} |
| 14 | |
| 15 | def __call__(self, im, tnf_type, theta=None): |
| 16 | ''' |
| 17 | Input: im: Torch.tensor[C, H, W] |
| 18 | tnf_type: 'random', 'affine', 'hom', 'tps' |
| 19 | Output: im: Torch.tensor[C, H, W] |
| 20 | grid: Torch.tensor[H, W, 2] |
| 21 | ''' |
| 22 | if tnf_type == 'random': |
| 23 | tnf_type = self.tnf_list[np.random.randint(3)] |
| 24 | theta = self.theta_generator[tnf_type]() |
| 25 | batch = {'image': im[None], 'theta': theta[None]} |
| 26 | else: |
| 27 | batch = {'image': im[None], 'theta': theta[None]} |
| 28 | tnf_res = self.tnf_generators[tnf_type](batch) |
| 29 | im = tnf_res['target_image'][0] |
| 30 | grid = refine_grid(tnf_res['warped_grid'][0]) |
| 31 | |
| 32 | return im, grid, tnf_type, theta |
| 33 |
nothing calls this directly
no test coverage detected