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)
