def op_script(data: torch.Tensor) -> torch.Tensor:
return kornia.hsv_to_rgb(data)
data = torch.tensor([[[[21., 22.], [22., 22.]],
[[13., 14.], [14., 14.]],
[[8., 8.], [8., 8.]]]]) # 3x2x2
actual = op_script(data)
expected = kornia.hsv_to_rgb(data)
assert_allclose(actual, expected)
def augment(im, aug_type):
''' Augment images as per given augmentation type. '''
if aug_type == 'normal':
im = im
elif aug_type == 'rotated':
rot_angle = np.random.choice([-90, 90, 180])
im = kornia_affine(im, rot_angle, 'rotate')
elif aug_type == 'hsv':
adjustFactor = np.random.choice(
[-0.1, -0.075, -0.05, 0.05, 0.075, 0.1])
hue, sat, value = torch.chunk(kornia.rgb_to_hsv(im), chunks=3, dim=-3)
adjust_mat = (torch.ones(1, 512, 512) * adjustFactor).cuda()
hueNew = hue + hue * adjust_mat.cuda()
hueNew = torch.clamp(hueNew, -2 * np.pi, 2 * np.pi)
satNew = sat + sat * adjust_mat.cuda()
satNew = torch.clamp(satNew, -2 * np.pi, 2 * np.pi)
new_im = torch.cat([hueNew, satNew, value], dim=-3)
im = kornia.hsv_to_rgb(new_im)
elif aug_type == 'gauss_noise':
im = augment_gaussian_noise(im)
elif aug_type == 'mirror':
im = torch.flip(im, [-2])
return im
def _hsv_to_rgb(img: torch.Tensor) -> torch.Tensor:
return kornia.hsv_to_rgb(img)