Generate a synthetically warped training pair using an affine transformation.
| 145 | |
| 146 | |
| 147 | class SynthPairTnf(object): |
| 148 | """ |
| 149 | Generate a synthetically warped training pair using an affine transformation. |
| 150 | """ |
| 151 | |
| 152 | def __init__(self, use_cuda=True, geometric_model='affine', crop_factor=9 / 16, output_size=(240, 240), |
| 153 | padding_factor=0.5): |
| 154 | assert isinstance(use_cuda, (bool)) |
| 155 | assert isinstance(crop_factor, (float)) |
| 156 | assert isinstance(output_size, (tuple)) |
| 157 | assert isinstance(padding_factor, (float)) |
| 158 | self.use_cuda = use_cuda |
| 159 | self.crop_factor = crop_factor |
| 160 | self.padding_factor = padding_factor |
| 161 | self.out_h, self.out_w = output_size |
| 162 | self.rescalingTnf = GeometricTnf('affine', out_h=self.out_h, out_w=self.out_w, |
| 163 | use_cuda=self.use_cuda) |
| 164 | self.geometricTnf = GeometricTnf(geometric_model, out_h=self.out_h, out_w=self.out_w, |
| 165 | use_cuda=self.use_cuda) |
| 166 | |
| 167 | def __call__(self, batch): |
| 168 | image_batch, theta_batch = batch['image'], batch['theta'] |
| 169 | if self.use_cuda: |
| 170 | image_batch = image_batch.cuda() |
| 171 | theta_batch = theta_batch.cuda() |
| 172 | |
| 173 | b, c, h, w = image_batch.size() |
| 174 | |
| 175 | # generate symmetrically padded image for bigger sampling region |
| 176 | # image_batch = self.symmetricImagePad(image_batch, self.padding_factor) |
| 177 | image_batch = self.expandImagePad(image_batch, self.padding_factor) |
| 178 | |
| 179 | |
| 180 | # convert to variables |
| 181 | image_batch = Variable(image_batch, requires_grad=False) |
| 182 | theta_batch = Variable(theta_batch, requires_grad=False) |
| 183 | |
| 184 | # get cropped image |
| 185 | cropped_image_batch, cropped_grid = self.rescalingTnf(image_batch=image_batch, |
| 186 | theta_batch=None, |
| 187 | padding_factor=self.padding_factor, |
| 188 | crop_factor=self.crop_factor, |
| 189 | return_sampling_grid=True) # Identity is used as no theta given |
| 190 | |
| 191 | # get transformed image |
| 192 | warped_image_batch, warped_grid = self.geometricTnf(image_batch=image_batch, |
| 193 | theta_batch=theta_batch, |
| 194 | padding_factor=self.padding_factor, |
| 195 | crop_factor=self.crop_factor, |
| 196 | return_sampling_grid=True) # Identity is used as no theta given |
| 197 | |
| 198 | valid_mask = (warped_grid[:,:,0] >= 0) & (warped_grid[:,:,0] < w) & (warped_grid[:,:,1] >= 0) & (warped_grid[:,:,1] < h) |
| 199 | |
| 200 | return {'source_image': cropped_image_batch, |
| 201 | 'target_image': warped_image_batch, |
| 202 | 'cropped_grid': cropped_grid, |
| 203 | 'warped_grid': warped_grid, |
| 204 | 'valid_mask': valid_mask} |