def test_affine_translate(self, device): torch.manual_seed(0) translation = torch.rand(1, 2, device=device) * 2.0 input = torch.rand(1, 2, 3, 4, device=device) transform = kornia.Affine(translation=translation).to(device) actual = transform(input) expected = kornia.translate(input, translation) assert_allclose(actual, expected)
def make_training_batch(input_tensor, patch, patch_mask): # determine patch size H, W = PA_cfg.image_shape[-2:] PATCH_SIZE = int(np.floor(np.sqrt((H*W*PA_cfg.percentage)))) translate_space = [H-PATCH_SIZE+1, W-PATCH_SIZE+1] bs = input_tensor.size(0) training_batch = [] for b in range(bs): # random translation u_t = np.random.randint(low=0, high=translate_space[0]) v_t = np.random.randint(low=0, high=translate_space[1]) # random scaling and rotation scale = np.random.rand() * (PA_cfg.scale_max - PA_cfg.scale_min) + PA_cfg.scale_min scale = torch.Tensor([scale]) angle = np.random.rand() * (PA_cfg.rotate_max - PA_cfg.rotate_min) + PA_cfg.rotate_min angle = torch.Tensor([angle]) center = torch.Tensor([u_t+PATCH_SIZE/2, v_t+PATCH_SIZE/2]).unsqueeze(0) rotation_m = kornia.get_rotation_matrix2d(center, angle, scale) # warp three tensors temp_mask = patch_mask.unsqueeze(0) temp_input = input_tensor[b].unsqueeze(0) temp_patch = patch.unsqueeze(0) temp_mask = kornia.translate(temp_mask.float(), translation=torch.Tensor([u_t, v_t]).unsqueeze(0)) temp_patch = kornia.translate(temp_patch.float(), translation=torch.Tensor([u_t, v_t]).unsqueeze(0)) mask_warpped = kornia.warp_affine(temp_mask.float(), rotation_m, temp_mask.size()[-2:]) patch_warpped = kornia.warp_affine(temp_patch.float(), rotation_m, temp_patch.size()[-2:]) # overlay overlay = temp_input * (1 - mask_warpped) + patch_warpped * mask_warpped training_batch.append(overlay) training_batch = torch.cat(training_batch, dim=0) return training_batch
def test_affine_translate(self, device, dtype): # TODO: Remove when #666 is implemented if device.type == 'cuda': pytest.skip("Currently breaks in CUDA." "See https://github.com/kornia/kornia/issues/666") torch.manual_seed(0) translation = torch.rand(1, 2, device=device, dtype=dtype) * 2.0 input = torch.rand(1, 2, 3, 4, device=device, dtype=dtype) transform = kornia.Affine(translation=translation).to(device=device, dtype=dtype) actual = transform(input) expected = kornia.translate(input, translation) assert_close(actual, expected, atol=1e-4, rtol=1e-4)
def _transform(self, img: torch.Tensor, param: torch.Tensor) -> torch.Tensor: B, _, H, _ = img.shape translation = param * H return kornia.translate(img, translation.expand(B, 2))
def TranslateY(x, v): batch_size = v.size(0) translation = torch.zeros((batch_size, 2), device=x.device) translation[:, 1] = v return kornia.translate(x, translation)
def translate_y(img: torch.Tensor, mag: torch.Tensor) -> torch.Tensor: mag = torch.stack([torch.zeros_like(mag), mag * img.size(-2)], dim=1) return kornia.translate(img, mag)