def test_cardinality(self, device, dtype, batch_size, height, width):
img = torch.ones(batch_size,
3,
height,
width,
device=device,
dtype=dtype)
assert kornia.rgb_to_grayscale(img).shape == (batch_size, 1, height,
width)
def closure():
nonlocal global_step
global_step += 1
if torch.is_grad_enabled():
optimizer.zero_grad()
_, _, model_images = flame()
model_grays = rgb_to_grayscale(model_images)
model_grays = normalize(model_grays, model_grays.mean(),
model_grays.std())
model_heatmaps = estimator(model_grays)[-1]
# loss_landmarks = 50 * criterion_landmarks(get_target_landmarks(target_landmarks), original_landmarks)
# loss = loss_landmarks
# log('loss_landmarks', loss_landmarks.item(), global_step)
loss_simple = criterion_simple(model_heatmaps, heatmaps)
loss = loss_simple
log('loss_simple', loss_simple.item(), global_step)
model_heatmaps_gray = get_heatmap_gray(model_heatmaps)
if criterion_gp is not None:
loss_gp = arg.loss_gp_lambda * criterion_gp(
model_heatmaps_gray, heatmaps_gray) + 100
loss = loss + loss_gp
log('loss_gp', loss_gp.item(), global_step)
log('loss', loss.item(), global_step)
if loss.requires_grad:
loss.backward(retain_graph=True)
log_img(
'original',
derescale_0_1(heatmaps_gray[0].unsqueeze(0), 0,
255).detach().to(dtype=torch.uint8), global_step)
log_img(
'rendered',
derescale_0_1(model_images, 0,
255)[0].detach().to(dtype=torch.uint8), global_step)
log_img(
'target',
derescale_0_1(model_heatmaps_gray[0].unsqueeze(0), 0,
255).detach().to(dtype=torch.uint8), global_step)
return loss
def test_opencv(self, device, dtype):
data = torch.tensor(
[
[
[0.3944633, 0.8597369, 0.1670904, 0.2825457, 0.0953912],
[0.1251704, 0.8020709, 0.8933256, 0.9170977, 0.1497008],
[0.2711633, 0.1111478, 0.0783281, 0.2771807, 0.5487481],
[0.0086008, 0.8288748, 0.9647092, 0.8922020, 0.7614344],
[0.2898048, 0.1282895, 0.7621747, 0.5657831, 0.9918593],
],
[
[0.5414237, 0.9962701, 0.8947155, 0.5900949, 0.9483274],
[0.0468036, 0.3933847, 0.8046577, 0.3640994, 0.0632100],
[0.6171775, 0.8624780, 0.4126036, 0.7600935, 0.7279997],
[0.4237089, 0.5365476, 0.5591233, 0.1523191, 0.1382165],
[0.8932794, 0.8517839, 0.7152701, 0.8983801, 0.5905426],
],
[
[0.2869580, 0.4700376, 0.2743714, 0.8135023, 0.2229074],
[0.9306560, 0.3734594, 0.4566821, 0.7599275, 0.7557513],
[0.7415742, 0.6115875, 0.3317572, 0.0379378, 0.1315770],
[0.8692724, 0.0809556, 0.7767404, 0.8742208, 0.1522012],
[0.7708948, 0.4509611, 0.0481175, 0.2358997, 0.6900532],
],
],
device=device,
dtype=dtype,
)
# Output data generated with OpenCV 4.1.1: cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
expected = torch.tensor(
[[
[0.4684734, 0.8954562, 0.6064363, 0.5236061, 0.6106016],
[0.1709944, 0.5133104, 0.7915002, 0.5745703, 0.1680204],
[0.5279005, 0.6092287, 0.3034387, 0.5333768, 0.6064113],
[0.3503858, 0.5720159, 0.7052018, 0.4558409, 0.3261529],
[0.6988886, 0.5897652, 0.6532392, 0.7234108, 0.7218805],
]],
device=device,
dtype=dtype,
)
img_gray = kornia.rgb_to_grayscale(data)
assert_allclose(img_gray, expected)
def extract_features(self, im):
kpts = self.det.detect(im, None)
# We will not train anything, so let's save time and memory by no_grad()
with torch.no_grad():
timg = K.image_to_tensor(im, False).float() / 255.
timg = timg.to(self.device)
if timg.shape[1] == 3:
timg_gray = K.rgb_to_grayscale(timg)
else:
timg_gray = timg
# kornia expects keypoints in the local affine frame format.
# Luckily, kornia_moons has a conversion function
lafs = laf_from_opencv_SIFT_kpts(kpts, device=self.device)
lafs_new = self.aff(lafs, timg_gray)
patches = KF.extract_patches_from_pyramid(timg_gray, lafs_new, 32)
B, N, CH, H, W = patches.size()
# Descriptor accepts standard tensor [B, CH, H, W], while patches are [B, N, CH, H, W] shape
# So we need to reshape a bit :)
descs = self.desc(patches.view(B * N, CH, H,
W)).view(B * N,
-1).detach().cpu().numpy()
kpts = np.array([[kp.pt[0], kp.pt[1]] for kp in kpts])
return kpts, descs
def test_rgb_to_grayscale_batch(self, device):
batch_size, channels, height, width = 2, 3, 4, 5
img = torch.ones(batch_size, channels, height, width).to(device)
assert kornia.rgb_to_grayscale(img).shape == \
(batch_size, 1, height, width)
def test_rgb_to_grayscale(self, device):
channels, height, width = 3, 4, 5
img = torch.ones(channels, height, width).to(device)
assert kornia.rgb_to_grayscale(img).shape == (1, height, width)
def test_smoke(self, device, dtype):
C, H, W = 3, 4, 5
img = torch.rand(C, H, W, device=device, dtype=dtype)
assert isinstance(kornia.rgb_to_grayscale(img), torch.Tensor)
#############################
# Create a batch of images
xb_bgr = torch.cat([x_bgr, hflip(x_bgr), vflip(x_bgr), rot180(x_bgr)])
imshow(xb_bgr)
#############################
# Convert BGR to RGB
xb_rgb = kornia.bgr_to_rgb(xb_bgr)
imshow(xb_rgb)
#############################
# Convert RGB to grayscale
# NOTE: image comes in torch.uint8, and kornia assumes floating point type
xb_gray = kornia.rgb_to_grayscale(xb_rgb.float() / 255.0)
imshow(xb_gray)
#############################
# Convert RGB to HSV
xb_hsv = kornia.rgb_to_hsv(xb_rgb.float() / 255.0)
imshow(xb_hsv[:, 2:3])
#############################
# Convert RGB to YUV
# NOTE: image comes in torch.uint8, and kornia assumes floating point type
yuv = kornia.rgb_to_yuv(xb_rgb.float() / 255.0)
y_channel = torchvision.utils.make_grid(yuv, nrow=2)[0, :, :]
plt.imshow(y_channel, cmap='gray', vmin=0, vmax=1) # Displaying only y channel
plt.axis('off')
plt.show()
def test_rgb_to_grayscale_batch(self):
batch_size, channels, height, width = 2, 3, 4, 5
img = torch.ones(batch_size, channels, height, width)
assert K.rgb_to_grayscale(img).shape == \
(batch_size, 1, height, width)
def test_rgb_to_grayscale(self):
channels, height, width = 3, 4, 5
img = torch.ones(channels, height, width)
assert K.rgb_to_grayscale(img).shape == (1, height, width)
def main():
# load the images
BASE_IMAGE_URL1: str = "https://raw.githubusercontent.com/kornia/data/main/panda.jpg" # augmentation
BASE_IMAGE_URL2: str = "https://raw.githubusercontent.com/kornia/data/main/simba.png" # color
BASE_IMAGE_URL3: str = "https://raw.githubusercontent.com/kornia/data/main/girona.png" # enhance
BASE_IMAGE_URL4: str = "https://raw.githubusercontent.com/kornia/data/main/baby_giraffe.png" # morphology
BASE_IMAGE_URL5: str = "https://raw.githubusercontent.com/kornia/data/main/persistencia_memoria.jpg" # filters
BASE_IMAGE_URL6: str = "https://raw.githubusercontent.com/kornia/data/main/delorean.png" # geometry
OUTPUT_PATH = Path(__file__).absolute().parent / "source/_static/img"
os.makedirs(OUTPUT_PATH, exist_ok=True)
print(f"Pointing images to path {OUTPUT_PATH}.")
img1 = read_img_from_url(BASE_IMAGE_URL1)
img2 = read_img_from_url(BASE_IMAGE_URL2, img1.shape[-2:])
img3 = read_img_from_url(BASE_IMAGE_URL3, img1.shape[-2:])
img4 = read_img_from_url(BASE_IMAGE_URL4)
img5 = read_img_from_url(BASE_IMAGE_URL5, (234, 320))
img6 = read_img_from_url(BASE_IMAGE_URL6)
# TODO: make this more generic for modules out of kornia.augmentation
# Dictionary containing the transforms to generate the sample images:
# Key: Name of the transform class.
# Value: (parameters, num_samples, seed)
mod = importlib.import_module("kornia.augmentation")
augmentations_list: dict = {
"CenterCrop": ((184, 184), 1, 2018),
"ColorJitter": ((0.3, 0.3, 0.3, 0.3), 2, 2018),
"RandomAffine": (((-15.0, 20.0), (0.1, 0.1), (0.7, 1.3), 20), 2, 2019),
"RandomBoxBlur": (((7, 7), ), 1, 2020),
"RandomCrop": ((img1.shape[-2:], (50, 50)), 2, 2020),
"RandomChannelShuffle": ((), 1, 2020),
"RandomElasticTransform": (((63, 63), (32, 32), (2.0, 2.0)), 2, 2018),
"RandomEqualize": ((), 1, 2020),
"RandomErasing": (((0.2, 0.4), (0.3, 1 / 0.3)), 2, 2017),
"RandomFisheye": ((torch.tensor([-0.3, 0.3]), torch.tensor([-0.3,
0.3]),
torch.tensor([0.9, 1.0])), 2, 2020),
"RandomGaussianBlur": (((3, 3), (0.1, 2.0)), 1, 2020),
"RandomGaussianNoise": ((0.0, 0.05), 1, 2020),
"RandomGrayscale": ((), 1, 2020),
"RandomHorizontalFlip": ((), 1, 2020),
"RandomInvert": ((), 1, 2020),
"RandomMotionBlur": ((7, 35.0, 0.5), 2, 2020),
"RandomPerspective": ((0.2, ), 2, 2020),
"RandomPosterize": (((1, 4), ), 2, 2016),
"RandomResizedCrop":
((img1.shape[-2:], (1.0, 2.0), (1.0, 2.0)), 2, 2020),
"RandomRotation": ((45.0, ), 2, 2019),
"RandomSharpness": ((16.0, ), 1, 2019),
"RandomSolarize": ((0.2, 0.2), 2, 2019),
"RandomVerticalFlip": ((), 1, 2020),
"RandomThinPlateSpline": ((), 1, 2020),
}
# ITERATE OVER THE TRANSFORMS
for aug_name, (args, num_samples, seed) in augmentations_list.items():
img_in = img1.repeat(num_samples, 1, 1, 1)
# dynamically create the class instance
cls = getattr(mod, aug_name)
aug = cls(*args, p=1.0)
# set seed
torch.manual_seed(seed)
# apply the augmentaiton to the image and concat
out = aug(img_in)
if aug_name == "CenterCrop":
h, w = img1.shape[-2:]
h_new, w_new = out.shape[-2:]
h_dif, w_dif = int(h - h_new), int(w - w_new)
out = torch.nn.functional.pad(out,
(w_dif // 2, w_dif // 2, 0, h_dif))
out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))],
dim=-1)
# save the output image
out_np = K.utils.tensor_to_image((out * 255.0).byte())
cv2.imwrite(str(OUTPUT_PATH / f"{aug_name}.png"), out_np)
sig = f"{aug_name}({', '.join([str(a) for a in args])}, p=1.0)"
print(f"Generated image example for {aug_name}. {sig}")
mod = importlib.import_module("kornia.augmentation")
mix_augmentations_list: dict = {
"RandomMixUp": (((0.3, 0.4), ), 2, 20),
"RandomCutMix": ((img1.shape[-2], img1.shape[-1]), 2, 2019),
}
# ITERATE OVER THE TRANSFORMS
for aug_name, (args, num_samples, seed) in mix_augmentations_list.items():
img_in = torch.cat([img1, img2])
# dynamically create the class instance
cls = getattr(mod, aug_name)
aug = cls(*args, p=1.0)
# set seed
torch.manual_seed(seed)
# apply the augmentaiton to the image and concat
out, _ = aug(img_in, torch.tensor([0, 1]))
out = torch.cat(
[img_in[0], img_in[1], *(out[i] for i in range(out.size(0)))],
dim=-1)
# save the output image
out_np = K.utils.tensor_to_image((out * 255.0).byte())
cv2.imwrite(str(OUTPUT_PATH / f"{aug_name}.png"), out_np)
sig = f"{aug_name}({', '.join([str(a) for a in args])}, p=1.0)"
print(f"Generated image example for {aug_name}. {sig}")
mod = importlib.import_module("kornia.color")
color_transforms_list: dict = {
"grayscale_to_rgb": ((), 3),
"rgb_to_bgr": ((), 1),
"rgb_to_grayscale": ((), 1),
"rgb_to_hsv": ((), 1),
"rgb_to_hls": ((), 1),
"rgb_to_luv": ((), 1),
"rgb_to_lab": ((), 1),
# "rgb_to_rgba": ((1.,), 1),
"rgb_to_xyz": ((), 1),
"rgb_to_ycbcr": ((), 1),
"rgb_to_yuv": ((), 1),
"rgb_to_linear_rgb": ((), 1),
}
# ITERATE OVER THE TRANSFORMS
for fn_name, (args, num_samples) in color_transforms_list.items():
# import function and apply
fn = getattr(mod, fn_name)
if fn_name == "grayscale_to_rgb":
out = fn(K.rgb_to_grayscale(img2), *args)
else:
out = fn(img2, *args)
# perform normalization to visualize
if fn_name == "rgb_to_lab":
out = out[:, :1] / 100.0
elif fn_name == "rgb_to_hsv":
out[:, :1] = out[:, :1] / 2 * math.pi
elif fn_name == "rgb_to_luv":
out = out[:, :1] / 116.0
# repeat channels for grayscale
if out.shape[1] != 3:
out = out.repeat(1, 3, 1, 1)
# save the output image
if fn_name == "grayscale_to_rgb":
out = torch.cat([
K.rgb_to_grayscale(img2[0]).repeat(3, 1, 1),
*(out[i] for i in range(out.size(0)))
],
dim=-1)
else:
out = torch.cat([img2[0], *(out[i] for i in range(out.size(0)))],
dim=-1)
out_np = K.utils.tensor_to_image((out * 255.0).byte())
cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
sig = f"{fn_name}({', '.join([str(a) for a in args])})"
print(f"Generated image example for {fn_name}. {sig}")
# korna.enhance module
mod = importlib.import_module("kornia.enhance")
transforms: dict = {
"adjust_brightness": ((torch.tensor([0.25, 0.5]), ), 2),
"adjust_contrast": ((torch.tensor([0.65, 0.5]), ), 2),
"adjust_gamma": ((torch.tensor([0.85, 0.75]), 2.0), 2),
"adjust_hue": ((torch.tensor([-math.pi / 4, math.pi / 4]), ), 2),
"adjust_saturation": ((torch.tensor([1.0, 2.0]), ), 2),
"solarize": ((torch.tensor([0.8, 0.5]), torch.tensor([-0.25,
0.25])), 2),
"posterize": ((torch.tensor([4, 2]), ), 2),
"sharpness": ((torch.tensor([1.0, 2.5]), ), 2),
"equalize": ((), 1),
"invert": ((), 1),
"equalize_clahe": ((), 1),
"add_weighted": ((0.75, 0.25, 2.0), 1),
}
# ITERATE OVER THE TRANSFORMS
for fn_name, (args, num_samples) in transforms.items():
img_in = img3.repeat(num_samples, 1, 1, 1)
if fn_name == "add_weighted":
args_in = (img_in, args[0], img2, args[1], args[2])
else:
args_in = (img_in, *args)
# import function and apply
fn = getattr(mod, fn_name)
out = fn(*args_in)
# save the output image
out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))],
dim=-1)
out_np = K.utils.tensor_to_image((out * 255.0).byte())
cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
sig = f"{fn_name}({', '.join([str(a) for a in args])})"
print(f"Generated image example for {fn_name}. {sig}")
# korna.morphology module
mod = importlib.import_module("kornia.morphology")
kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
transforms: dict = {
"dilation": ((kernel, ), 1),
"erosion": ((kernel, ), 1),
"opening": ((kernel, ), 1),
"closing": ((kernel, ), 1),
"gradient": ((kernel, ), 1),
"top_hat": ((kernel, ), 1),
"bottom_hat": ((kernel, ), 1),
}
# ITERATE OVER THE TRANSFORMS
for fn_name, (args, num_samples) in transforms.items():
img_in = img4.repeat(num_samples, 1, 1, 1)
args_in = (img_in, *args)
# import function and apply
# import pdb;pdb.set_trace()
fn = getattr(mod, fn_name)
out = fn(*args_in)
# save the output image
out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))],
dim=-1)
out_np = K.utils.tensor_to_image((out * 255.0).byte())
cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
sig = f"{fn_name}({', '.join([str(a) for a in args])})"
print(f"Generated image example for {fn_name}. {sig}")
# korna.filters module
mod = importlib.import_module("kornia.filters")
kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
transforms: dict = {
"box_blur": (((5, 5), ), 1),
"median_blur": (((5, 5), ), 1),
"gaussian_blur2d": (((5, 5), (1.5, 1.5)), 1),
"motion_blur": ((5, 90.0, 1.0), 1),
"max_blur_pool2d": ((5, ), 1),
"blur_pool2d": ((5, ), 1),
"unsharp_mask": (((5, 5), (1.5, 1.5)), 1),
"laplacian": ((5, ), 1),
"sobel": ((), 1),
"spatial_gradient": ((), 1),
"canny": ((), 1),
}
# ITERATE OVER THE TRANSFORMS
for fn_name, (args, num_samples) in transforms.items():
img_in = img5.repeat(num_samples, 1, 1, 1)
args_in = (img_in, *args)
# import function and apply
fn = getattr(mod, fn_name)
out = fn(*args_in)
if fn_name in ("max_blur_pool2d", "blur_pool2d"):
out = K.geometry.resize(out, img_in.shape[-2:])
if fn_name == "canny":
out = out[1].repeat(1, 3, 1, 1)
if isinstance(out, torch.Tensor):
out = out.clamp(min=0.0, max=1.0)
if fn_name in ("laplacian", "sobel", "spatial_gradient", "canny"):
out = K.enhance.normalize_min_max(out)
if fn_name == "spatial_gradient":
out = out.permute(2, 1, 0, 3, 4).squeeze()
# save the output image
out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))],
dim=-1)
out_np = K.utils.tensor_to_image((out * 255.0).byte())
cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
sig = f"{fn_name}({', '.join([str(a) for a in args])})"
print(f"Generated image example for {fn_name}. {sig}")
# korna.geometry.transform module
mod = importlib.import_module("kornia.geometry.transform")
h, w = img6.shape[-2:]
def _get_tps_args():
src = torch.tensor([[[-1.0, -1.0], [-1.0, 1.0], [1.0,
-1.0], [1.0, -1.0],
[0.0, 0.0]]]).repeat(2, 1, 1) # Bx5x2
dst = src + torch.distributions.Uniform(-0.2, 0.2).rsample((2, 5, 2))
kernel, affine = K.geometry.transform.get_tps_transform(dst, src)
return src, kernel, affine
transforms: dict = {
"warp_affine": (
(
K.geometry.transform.get_affine_matrix2d(
translations=torch.zeros(2, 2),
center=(torch.tensor([w, h]) / 2).repeat(2, 1),
scale=torch.distributions.Uniform(0.5, 1.5).rsample(
(2, 2)),
angle=torch.tensor([-25.0, 25.0]),
)[:, :2, :3],
(h, w),
),
2,
),
"remap": (
(
*(K.utils.create_meshgrid(h, w, normalized_coordinates=True) -
0.25).unbind(-1),
'bilinear',
'zeros',
True,
True,
),
1,
),
"warp_image_tps": ((_get_tps_args()), 2),
"rotate": ((torch.tensor([-15.0, 25.0]), ), 2),
"translate": ((torch.tensor([[10.0, -15], [50.0, -25.0]]), ), 2),
"scale": ((torch.tensor([[0.5, 1.25], [1.0, 1.5]]), ), 2),
"shear": ((torch.tensor([[0.1, -0.2], [-0.2, 0.1]]), ), 2),
"rot180": ((), 1),
"hflip": ((), 1),
"vflip": ((), 1),
"resize": (((120, 220), ), 1),
"rescale": ((0.5, ), 1),
"elastic_transform2d":
((torch.rand(1, 2, h, w) * 2 - 1, (63, 63), (32, 32), (4.0, 4.0)), 1),
"pyrdown": ((), 1),
"pyrup": ((), 1),
"build_pyramid": ((3, ), 1),
}
# ITERATE OVER THE TRANSFORMS
for fn_name, (args, num_samples) in transforms.items():
img_in = img6.repeat(num_samples, 1, 1, 1)
args_in = (img_in, *args)
# import function and apply
fn = getattr(mod, fn_name)
out = fn(*args_in)
if fn_name in ("resize", "rescale", "pyrdown", "pyrup"):
h_new, w_new = out.shape[-2:]
out = torch.nn.functional.pad(out,
(0, (w - w_new), 0, (h - h_new)))
if fn_name == "build_pyramid":
_out = []
for pyr in out[1:]:
h_new, w_new = pyr.shape[-2:]
out_tmp = torch.nn.functional.pad(pyr, (0, (w - w_new), 0,
(h - h_new)))
_out.append(out_tmp)
out = torch.cat(_out)
# save the output image
out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))],
dim=-1)
out_np = K.utils.tensor_to_image((out * 255.0).byte())
cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
sig = f"{fn_name}({', '.join([str(a) for a in args])})"
print(f"Generated image example for {fn_name}. {sig}")
