(self, geometric_model='affine', tps_grid_size=3, tps_reg_factor=0, out_h=240, out_w=240,
offset_factor=None, use_cuda=True)
| 78 | """ |
| 79 | |
| 80 | def __init__(self, geometric_model='affine', tps_grid_size=3, tps_reg_factor=0, out_h=240, out_w=240, |
| 81 | offset_factor=None, use_cuda=True): |
| 82 | self.out_h = out_h |
| 83 | self.out_w = out_w |
| 84 | self.geometric_model = geometric_model |
| 85 | self.use_cuda = use_cuda |
| 86 | self.offset_factor = offset_factor |
| 87 | |
| 88 | if geometric_model == 'affine' and offset_factor is None: |
| 89 | self.gridGen = AffineGridGen(out_h=out_h, out_w=out_w, use_cuda=use_cuda) |
| 90 | elif geometric_model == 'affine' and offset_factor is not None: |
| 91 | self.gridGen = AffineGridGenV2(out_h=out_h, out_w=out_w, use_cuda=use_cuda) |
| 92 | elif geometric_model == 'hom': |
| 93 | self.gridGen = HomographyGridGen(out_h=out_h, out_w=out_w, use_cuda=use_cuda) |
| 94 | elif geometric_model == 'tps': |
| 95 | self.gridGen = TpsGridGen(out_h=out_h, out_w=out_w, grid_size=tps_grid_size, |
| 96 | reg_factor=tps_reg_factor, use_cuda=use_cuda) |
| 97 | if offset_factor is not None: |
| 98 | self.gridGen.grid_X = self.gridGen.grid_X / offset_factor |
| 99 | self.gridGen.grid_Y = self.gridGen.grid_Y / offset_factor |
| 100 | |
| 101 | self.theta_identity = torch.Tensor(np.expand_dims(np.array([[1, 0, 0], [0, 1, 0]]), 0).astype(np.float32)) |
| 102 | if use_cuda: |
| 103 | self.theta_identity = self.theta_identity.cuda() |
| 104 | |
| 105 | def __call__(self, image_batch, theta_batch=None, out_h=None, out_w=None, return_warped_image=True, |
| 106 | return_sampling_grid=False, padding_factor=1.0, crop_factor=1.0): |
nothing calls this directly
no test coverage detected