Geometric transfromation to an image batch (wrapped in a PyTorch Variable) ( can be used with no transformation to perform bilinear resizing )
| 72 | |
| 73 | |
| 74 | class GeometricTnf(object): |
| 75 | """ |
| 76 | Geometric transfromation to an image batch (wrapped in a PyTorch Variable) |
| 77 | ( can be used with no transformation to perform bilinear resizing ) |
| 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): |
| 107 | |
| 108 | crop_factor = random.uniform(0.5, 5) |
| 109 | if image_batch is None: |
| 110 | b = 1 |
| 111 | else: |
| 112 | b = image_batch.size(0) |
| 113 | if theta_batch is None: |
| 114 | theta_batch = self.theta_identity |
| 115 | theta_batch = theta_batch.expand(b, 2, 3).contiguous() |
| 116 | theta_batch = Variable(theta_batch, requires_grad=False) |
| 117 | # check if output dimensions have been specified at call time and have changed |
| 118 | if (out_h is not None and out_w is not None) and (out_h != self.out_h or out_w != self.out_w): |
| 119 | if self.geometric_model == 'affine': |
| 120 | gridGen = AffineGridGen(out_h, out_w, use_cuda=self.use_cuda) |
| 121 | elif self.geometric_model == 'hom': |
| 122 | gridGen = HomographyGridGen(out_h, out_w, use_cuda=self.use_cuda) |
| 123 | elif self.geometric_model == 'tps': |
| 124 | gridGen = TpsGridGen(out_h, out_w, use_cuda=self.use_cuda) |
| 125 | else: |
| 126 | gridGen = self.gridGen |
| 127 | sampling_grid = gridGen(theta_batch) |
| 128 | |
| 129 | # rescale grid according to crop_factor and padding_factor |
| 130 | if padding_factor != 1 or crop_factor != 1: |
| 131 | sampling_grid = sampling_grid * (padding_factor * crop_factor) |