(self, batch)
| 342 | use_cuda=self.use_cuda) |
| 343 | |
| 344 | def __call__(self, batch): |
| 345 | image_batch, theta_batch = batch['image'], batch['theta'] |
| 346 | # theta_aff=torch.index_select(theta_batch[:,:6],1,self.aff_reorder_idx) |
| 347 | theta_aff = theta_batch[:, :6].contiguous() |
| 348 | theta_tps = theta_batch[:, 6:] |
| 349 | |
| 350 | if self.use_cuda: |
| 351 | image_batch = image_batch.cuda() |
| 352 | theta_aff = theta_aff.cuda() |
| 353 | theta_tps = theta_tps.cuda() |
| 354 | |
| 355 | b, c, h, w = image_batch.size() |
| 356 | |
| 357 | # generate symmetrically padded image for bigger sampling region |
| 358 | image_batch = self.symmetricImagePad(image_batch, self.padding_factor) |
| 359 | |
| 360 | # convert to variables |
| 361 | image_batch = Variable(image_batch, requires_grad=False) |
| 362 | theta_aff = Variable(theta_aff, requires_grad=False) |
| 363 | theta_tps = Variable(theta_tps, requires_grad=False) |
| 364 | |
| 365 | # get cropped image |
| 366 | cropped_image_batch = self.rescalingTnf(image_batch=image_batch, |
| 367 | theta_batch=None, |
| 368 | padding_factor=self.padding_factor, |
| 369 | crop_factor=self.crop_factor) # Identity is used as no theta given |
| 370 | # get transformed image |
| 371 | warped_image_aff = self.affTnf(image_batch=image_batch, |
| 372 | theta_batch=theta_aff, |
| 373 | padding_factor=self.padding_factor, |
| 374 | crop_factor=self.crop_factor) |
| 375 | |
| 376 | warped_image_tps = self.tpsTnf(image_batch=image_batch, |
| 377 | theta_batch=theta_tps, |
| 378 | padding_factor=self.padding_factor, |
| 379 | crop_factor=self.crop_factor) |
| 380 | |
| 381 | return {'source_image': cropped_image_batch, 'target_image_aff': warped_image_aff, |
| 382 | 'target_image_tps': warped_image_tps, 'theta_GT_aff': theta_aff, 'theta_GT_tps': theta_tps} |
| 383 | |
| 384 | |
| 385 | class AffineGridGen(Module): |
nothing calls this directly
no test coverage detected