Files
kornia--kornia/docs/generate_examples.py
wehub-resource-sync 3a2c66702c
Lint / TOML Format (push) Waiting to run
Tests on CPU (scheduled) / check-skip (push) Waiting to run
Tests on CPU (scheduled) / pre-tests (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-ubuntu (float32) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-ubuntu (float64) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.11, float32, 2.5.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.11, float32, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.11, float64, 2.5.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.11, float64, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.12, float32, 2.5.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.12, float32, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.12, float64, 2.5.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.12, float64, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.13, float32, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-windows (3.13, float64, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-mac (3.11, float32, 2.5.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-mac (3.11, float32, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-mac (3.12, float32, 2.5.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-mac (3.12, float32, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / tests-cpu-mac (3.13, float32, 2.9.1) (push) Blocked by required conditions
Tests on CPU (scheduled) / coverage (push) Blocked by required conditions
Tests on CPU (scheduled) / typing (push) Blocked by required conditions
Tests on CPU (scheduled) / tutorials (push) Blocked by required conditions
Tests on CPU (scheduled) / docs (push) Blocked by required conditions
chore: import upstream snapshot with attribution
2026-07-13 12:49:27 +08:00

752 lines
32 KiB
Python

# LICENSE HEADER MANAGED BY add-license-header
#
# Copyright 2018 Kornia Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from __future__ import annotations
import importlib
import math
import os
from pathlib import Path
from typing import Optional
import cv2
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import requests
import torch
import torch.nn.functional as F
from kornia_moons.feature import visualize_LAF
import kornia as K
mpl.use("Agg")
def download_tutorials_examples(download_infos: dict[str, str], directory: Path):
URL_BASE = "https://raw.githubusercontent.com/kornia/tutorials/master/"
for filename, path in download_infos.items():
url = URL_BASE + path
# perform request
response = requests.get(url, timeout=60).content
path = directory / filename
with open(path, "wb") as fp:
fp.write(response)
def read_img_from_url(url: str, resize_to: Optional[tuple[int, int]] = None, **resize_kwargs) -> torch.Tensor:
# perform request
response = requests.get(url, timeout=60).content
# convert to array of ints
nparr = np.frombuffer(response, np.uint8)
# convert to image array and resize
img: np.ndarray = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)[..., :3]
# convert the image to a tensor
img_t: torch.Tensor = K.image.image_to_tensor(img, keepdim=False) # 1xCxHXW
img_t = img_t.float() / 255.0
if resize_to is None:
img_t = K.geometry.resize(img_t, 184, **resize_kwargs)
else:
img_t = K.geometry.resize(img_t, resize_to, **resize_kwargs)
return img_t
def transparent_pad(src: torch.Tensor, shape: tuple[int, int]) -> torch.Tensor:
"""Apply a transparent pad to src (centerized) to match with shape (h, w)"""
w_pad = abs(int(src.shape[-1] - shape[-1]) // 2)
h_pad = abs(int(src.shape[-2] - shape[-2]) // 2)
return F.pad(K.color.rgb_to_rgba(src, 1.0), (w_pad, w_pad, h_pad, h_pad), "constant", 0.0)
def draw_bbox_kpts(imgs: torch.Tensor, bboxes: torch.Tensor, keypoints: torch.Tensor) -> torch.Tensor:
rectangle = torch.zeros(imgs.shape[0], imgs.shape[1], 4)
rectangle[..., 0] = bboxes[..., 0] # x1
rectangle[..., 1] = bboxes[..., 1] # y1
rectangle[..., 2] = bboxes[..., 0] + bboxes[..., -2] # x2
rectangle[..., 3] = bboxes[..., 1] + bboxes[..., -1] # y2
color = torch.tensor([1, 0, 0]).repeat(imgs.shape[0], imgs.shape[1], 1)
imgs_draw = K.image.draw_rectangle(imgs, rectangle, color=color)
rectangle2 = torch.zeros(imgs.shape[0], imgs.shape[1], 4)
for n in range(keypoints.shape[-2]):
rectangle2[..., n, 0] = keypoints[..., n, 0] - 2
rectangle2[..., n, 1] = keypoints[..., n, 1] - 2
rectangle2[..., n, 2] = keypoints[..., n, 0] + 2
rectangle2[..., n, 3] = keypoints[..., n, 1] + 2
color = torch.tensor([0, 0, 1]).repeat(imgs.shape[0], imgs.shape[1], 1)
imgs_draw = K.utils.draw_rectangle(imgs_draw, rectangle2, color=color, fill=True)
return imgs_draw
def main():
# Download the tutorial examples for the main docs
# Note: Training API examples (image_classifier, object_detection, semantic_segmentation) removed
# as they depend on kornia.x which has been removed
URLS_TUTORIALS_EXAMPLES = {}
OUTPUT_PATH_SCRIPTS = Path(__file__).absolute().parent / "source/_static/scripts/"
os.makedirs(OUTPUT_PATH_SCRIPTS, exist_ok=True)
print(f"Downloading script examples from kornia/tutorials. Saving into the path {OUTPUT_PATH_SCRIPTS}.")
download_tutorials_examples(URLS_TUTORIALS_EXAMPLES, OUTPUT_PATH_SCRIPTS)
# 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
hash1 = "8b98f44abbe92b7a84631ed06613b08fee7dae14"
BASE_IMAGEOUTDOOR_URL7: str = f"https://github.com/kornia/data_test/raw/{hash1}/knchurch_disk.pt" # image matching
BASE_IMAGEOUTDOOR_URL8: str = ( # Response functions
"https://github.com/kornia/data/raw/main/kornia_banner_pixie.png"
)
MASK_IMAGE_URL2: str = "https://raw.githubusercontent.com/kornia/data/main/simba_mask.png"
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)
img_kornia = read_img_from_url(BASE_IMAGEOUTDOOR_URL8)
# Read the masks as (B, H, W)
mask2 = read_img_from_url(MASK_IMAGE_URL2, img1.shape[-2:], interpolation="nearest")
mask2 = mask2.median(dim=1)[0]
# 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),
"ColorJiggle": ((0.3, 0.3, 0.3, 0.3), 2, 2018),
"ColorJitter": ((0.3, 0.3, 0.3, 0.3), 2, 2022),
"PadTo": (((220, 450),), 1, 2022),
"RandomAffine": (((-15.0, 20.0), (0.1, 0.1), (0.7, 1.3), 20), 2, 2019),
"RandomBoxBlur": (((7, 7),), 1, 2020),
"RandomBrightness": (((0.0, 1.0),), 2, 2022),
"RandomContrast": (((0.0, 1.0),), 2, 2022),
"RandomCrop": ((img1.shape[-2:], (50, 50)), 2, 2020),
"RandomChannelDropout": ((), 1, 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),
"RandomGamma": (((0.0, 1.0),), 2, 2022),
"RandomGaussianBlur": (((3, 3), (0.1, 2.0)), 1, 2020),
"RandomGaussianIllumination": (((0.5, 0.5), (0.5, 0.5), (0.5, 0.5), (-1.0, 1.0)), 2, 2021),
"RandomGaussianNoise": ((0.0, 0.05), 1, 2020),
"RandomGrayscale": ((), 1, 2020),
"RandomHue": (((-0.5, 0.5),), 2, 2022),
"RandomHorizontalFlip": ((), 1, 2020),
"RandomInvert": ((), 1, 2020),
"RandomJPEG": (((1.0, 5.0),), 1, 2024),
"RandomLinearCornerIllumination": (((0.5, 0.5), (-1.0, 1.0)), 2, 2021),
"RandomLinearIllumination": (((0.5, 0.5), (-1.0, 1.0)), 2, 2021),
"RandomMedianBlur": (((3, 3),), 1, 2023),
"RandomMotionBlur": ((7, 35.0, 0.5), 2, 2020),
"RandomPerspective": ((0.2,), 2, 2020),
"RandomPlanckianJitter": ((), 2, 2022),
"RandomPlasmaShadow": (((0.2, 0.5),), 2, 2022),
"RandomPlasmaBrightness": ((), 2, 2022),
"RandomPlasmaContrast": ((), 2, 2022),
"RandomPosterize": (((1, 4),), 2, 2016),
"RandomResizedCrop": ((img1.shape[-2:], (1.0, 2.0), (1.0, 2.0)), 2, 2020),
"RandomRotation": ((45.0,), 2, 2019),
"RandomSaltAndPepperNoise": (((0.05, 0.5), (0.1, 0.7)), 2, 2024),
"RandomSaturation": (((0.5, 5.0),), 2, 2022),
"RandomSharpness": ((16.0,), 1, 2019),
"RandomSolarize": ((0.2, 0.2), 2, 2019),
"RandomVerticalFlip": ((), 1, 2020),
"RandomThinPlateSpline": ((), 1, 2020),
"RandomJigsaw": ((), 2, 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)
try:
aug = cls(*args, p=1.0)
except TypeError:
aug = cls(*args)
# set seed
torch.manual_seed(seed)
if aug_name == "RandomJigsaw": # make sure the image is dividable
img_in = K.geometry.resize(img_in, (1020, 500))
elif aug_name == "RandomJPEG":
img_in = img_in[..., :176, :]
# apply the augmentation to the image and concat
out = aug(img_in)
# save ori image to concatenate into the out image
ori = img_in[0]
if aug_name == "CenterCrop":
# Convert to RGBA, and center the output image with transparent pad
out = transparent_pad(out, tuple(img1[-2:].shape))
ori = K.color.rgb_to_rgba(ori, 1.0) # To match the dims
elif aug_name == "PadTo":
# Convert to RGBA, and center the original image with transparent pad
ori = transparent_pad(img_in[0], tuple(out.shape[-2:]))
out = K.color.rgb_to_rgba(out, 1.0) # To match the dims
out = torch.cat([ori, *(out[i] for i in range(out.size(0)))], dim=-1)
# save the output image
out_np = K.image.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}")
mix_augmentations_list = {"RandomMixUpV2": ((), 2, 20), "RandomCutMixV2": ((), 2, 2019), "PatchMix": ((), 2, 2024)}
# ITERATE OVER THE TRANSFORMS
for aug_name, (args, _, 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 augmentation to the image and concat
# PatchMix returns (B, C, H, W); index [0] to get (C, H, W) for cat
if aug_name == "PatchMix":
img_aug = aug(img_in)[0]
else:
img_aug, _ = aug(img_in, torch.tensor([0, 1]))
output = torch.cat([img_in[0], img_in[1], img_aug], dim=-1)
# save the output image
out_np = K.utils.tensor_to_image((output * 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}")
mask_augmentations_list = {"RandomTransplantation": (([0],), 0)}
# ITERATE OVER THE TRANSFORMS
for aug_name, (args, seed) in mask_augmentations_list.items():
img_in = torch.cat([img1, img2])
mask_in = torch.cat([torch.zeros_like(mask2), mask2])
# dynamically create the class instance
cls = getattr(mod, aug_name)
aug = cls(*args, p=1.0)
# set seed
torch.manual_seed(seed)
# apply the augmentation to the image and concat
img_aug, _ = aug(img_in, mask_in)
output = torch.cat([img_in[0], img_in[1], img_aug[0]], dim=-1)
# save the output image
out_np = K.utils.tensor_to_image((output * 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}")
# Containers
aug_container_list = {
"AugmentationSequential": (
{
"args": (
K.augmentation.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0),
K.augmentation.RandomAffine(360, [0.1, 0.1], [0.7, 1.2], [30.0, 50.0], p=1.0),
K.augmentation.RandomPerspective(0.5, p=1.0),
),
"data_keys": ["input", "bbox_xywh", "keypoints"],
},
(
torch.tensor([[[125, 5, 115, 80]]], dtype=torch.float32), # bbox
torch.tensor([[[166, 42], [197, 42]]], dtype=torch.float32), # keypoints
),
2,
2023,
),
"PatchSequential": (
{
"args": (
K.augmentation.ColorJitter(0.2, 0.1, 0.1, 0.1, p=1),
K.augmentation.RandomAffine(10, [0.1, 0.2], [0.7, 1.2], [0.0, 15.0], p=1),
K.augmentation.RandomPerspective(0.3, p=1),
K.augmentation.RandomSolarize(0.01, 0.05, p=0.6),
),
"grid_size": (2, 2),
"same_on_batch": False,
"patchwise_apply": False,
},
(),
2,
2023,
),
}
for aug_name, (args, labels, num_samples, seed) in aug_container_list.items():
img_in = img1.repeat(num_samples, 1, 1, 1)
cls = getattr(mod, aug_name)
tfms = args.pop("args")
augs = cls(*tfms, **args)
# set seed
torch.manual_seed(seed)
if aug_name == "PatchSequential":
out = augs(img_in)
inp = img_in
else:
labels = (labels[0].expand(num_samples, -1, -1), labels[1].expand(num_samples, -1, -1))
out = augs(img_in, *labels)
out = draw_bbox_kpts(out[0], out[1].int(), out[2].int())
inp = draw_bbox_kpts(img_in, labels[0].int(), labels[1].int())
output = torch.cat([inp[0], *(out[i] for i in range(out.size(0)))], dim=-1)
# save the output image
out_np = K.utils.tensor_to_image((output * 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),
"apply_colormap": ((K.color.ColorMap("autumn", 256),), 1),
"ApplyColorMap": ((K.color.ColorMap("winter", 256),), 1),
}
# ITERATE OVER THE TRANSFORMS
for fn_name, (args, _) in color_transforms_list.items():
# import function and apply
fn = getattr(mod, fn_name)
if fn_name == "grayscale_to_rgb":
out = fn(K.color.rgb_to_grayscale(img2), *args)
elif fn_name == "apply_colormap":
gray_image = (K.color.rgb_to_grayscale(img2) * 255.0).round()
out = K.color.rgb_to_bgr(fn(gray_image, *args))
elif fn_name == "ApplyColorMap":
gray_image = (K.color.rgb_to_grayscale(img2) * 255.0).round()
out = K.color.rgb_to_bgr(fn(*args)(gray_image))
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 in ("grayscale_to_rgb", "apply_colormap", "ApplyColorMap"):
out = torch.cat(
[K.color.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.image.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.color.colormap
colormaps_list = {"AUTUMN": (256,)}
bar_img_gray = torch.arange(0, 256).repeat(1, 40, 1) # 1x1x40x256
bar_img = K.color.grayscale_to_rgb(bar_img_gray)
# ITERATE OVER THE COLORMAPS
for colormap_name, args in colormaps_list.items():
cm = K.color.ColorMap(base=colormap_name, num_colors=args[0])
out = K.color.rgb_to_bgr(K.color.apply_colormap(bar_img_gray, cm))[0]
out = torch.cat([bar_img, out], dim=-1)
out_np = K.image.tensor_to_image((out * 255.0).byte())
cv2.imwrite(str(OUTPUT_PATH / f"{colormap_name}.png"), out_np)
sig = f"{colormap_name}({', '.join([str(a) for a in args])})"
print(f"Generated image example for {colormap_name}. {sig}")
# Plot for all ColorMaps (ColorMapType.png)
height_image = 40
num_colors = 256
num_columns = 3
# 1 x height_image x num_colors
input_tensor = torch.arange(start=0, end=num_colors, step=1).unsqueeze(0).repeat(1, height_image, 1)
input_tensor = input_tensor.to("cpu").to(torch.float32)
# Get colormap list
colormap_list = K.color.ColorMapType.list()
num_colormaps = len(colormap_list)
# Calculate number of rows needed
num_rows = (num_colormaps + num_columns - 1) // num_columns
# Create figure and axis objects
fig, axes = plt.subplots(num_rows, num_columns, figsize=(12, 8))
for i, ax in enumerate(axes.flat):
if i < num_colormaps:
cmap = K.color.ColorMap(base=colormap_list[i], num_colors=num_colors)
res = K.color.ApplyColorMap(colormap=cmap)(input_tensor)[0]
ax.imshow(res.permute(1, 2, 0).numpy())
ax.set_title(colormap_list[i], fontsize=12)
ax.axis("off")
else:
fig.delaxes(ax)
fig.tight_layout()
fig.savefig(os.path.join(OUTPUT_PATH, "ColorMapType.png"), dpi=300)
# 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),
"jpeg_codec_differentiable": ((torch.tensor([50]),), 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 == "jpeg_codec_differentiable":
img_in = img_in[..., :176, :]
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.image.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.image.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 = {
"bilateral_blur": (((11, 11), 0.1, (3, 3)), 1),
"joint_bilateral_blur": (((11, 11), 0.1, (3, 3)), 1),
"box_blur": (((5, 5),), 1),
"median_blur": (((5, 5),), 1),
"gaussian_blur2d": (((5, 5), (1.5, 1.5)), 1),
"guided_blur": (((5, 5), 0.01), 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)
if fn_name == "joint_bilateral_blur":
guide = K.geometry.resize(img2.repeat(num_samples, 1, 1, 1), img_in.shape[-2:])
args_in = (img_in, guide, *args)
elif fn_name == "guided_blur":
args_in = (img_in, img_in, *args)
else:
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()
if fn_name == "joint_bilateral_blur":
out = torch.cat([args_in[1], out], dim=-1)
# save the output image
out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
out_np = K.image.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}")
# kornia.filters.in_range
mod = importlib.import_module("kornia.filters")
transforms: dict = {
"in_range": (((0.314, 0.2, 0.2), (0.47, 1.0, 1.0), True), 1),
}
# ITERATE OVER THE TRANSFORMS
for fn_name, (args, _) in transforms.items():
img_hsv = K.color.rgb_to_hsv(img1)
h, s, v = torch.split(img_hsv, split_size_or_sections=1, dim=1)
h = h / (2 * torch.pi)
img_hsv = torch.cat((h, s, v), dim=1)
args_in = (img_hsv, *args)
fn = getattr(mod, fn_name)
mask = fn(*args_in)
filtered = img1 * mask
mask = mask.repeat(1, img1.shape[1], 1, 1)
# save the output image
out = torch.cat([img1[0], mask[0], filtered[0]], dim=-1)
out_np = K.image.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.geometry.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),
"build_laplacian_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)
if fn_name == "build_laplacian_pyramid":
h_, w_ = out[0].shape[-2:]
_out = [out[0]]
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)))
print(out_tmp.size())
_out.append(out_tmp)
out = torch.cat(_out)
# save the output image
if fn_name != "build_laplacian_pyramid":
out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
else:
out = torch.cat([*(out[i] for i in range(out.size(0)))], dim=-1)
out_np = K.image.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}")
# Image Matching and local features
img_matching_data = torch.hub.load_state_dict_from_url(BASE_IMAGEOUTDOOR_URL7, map_location=torch.device("cpu"))
img_outdoor = img_matching_data["img2"]
print("Generating local feature detections ")
disk = K.feature.DISK.from_pretrained("depth")
with torch.no_grad():
disk_feat = disk(img_outdoor)[0]
xy = disk_feat.keypoints.detach().cpu().numpy()
cur_fname = str(OUTPUT_PATH / "disk_outdoor_depth.jpg")
plt.figure()
plt.imshow(K.tensor_to_image(img_outdoor))
plt.scatter(xy[:, 0], xy[:, 1], 3, color="lime")
plt.title('DISK("depth") keypoints')
plt.savefig(cur_fname)
plt.close()
kah = K.feature.KeyNetAffNetHardNet(512).eval()
with torch.no_grad():
lafs, _resps, _descs = kah(K.color.rgb_to_grayscale(img_outdoor))
fig1, ax = visualize_LAF(img_outdoor, lafs, color="lime", return_fig_ax=True, draw_ori=False)
ax.set_title("KeyNetAffNet 512 LAFs")
cur_fname = str(OUTPUT_PATH / "keynet_affnet.jpg")
fig1.savefig(cur_fname)
plt.close()
keynet = K.feature.KeyNetDetector(True, 512).eval()
with torch.no_grad():
lafs, _resps = keynet(K.color.rgb_to_grayscale(img_outdoor))
xy = K.feature.get_laf_center(lafs).detach().cpu().numpy().reshape(-1, 2)
cur_fname = str(OUTPUT_PATH / "keynet.jpg")
plt.figure()
plt.imshow(K.tensor_to_image(img_outdoor))
plt.scatter(xy[:, 0], xy[:, 1], 3, color="lime")
plt.title("KeyNet 512 keypoints")
plt.savefig(cur_fname)
plt.close()
# korna.feature module
mod = importlib.import_module("kornia.feature")
responses: list = [
"harris_response",
"gftt_response",
"hessian_response",
"dog_response_single",
"KeyNet",
"DISK",
"ALIKED",
"XFeat",
]
# ITERATE OVER THE TRANSFORMS
for fn_name in responses:
# import function and apply
img_in = K.color.rgb_to_grayscale(img_kornia)
if fn_name == "KeyNet":
fn = K.feature.KeyNet(True)
out = fn(img_in)
elif fn_name == "DISK":
fn = K.feature.DISK.from_pretrained("depth")
h, w = img_outdoor.shape[2:]
pd_h = 32 - h % 32 if h % 32 > 0 else 0
pd_w = 32 - w % 32 if w % 32 > 0 else 0
img_in = torch.nn.functional.pad(img_outdoor, (0, pd_w, 0, pd_h), value=0.0)
out, _ = fn.heatmap_and_dense_descriptors(img_in)
out = K.color.grayscale_to_rgb(out)
img_in = K.color.rgb_to_bgr(img_in)
elif fn_name == "ALIKED":
fn = K.feature.ALIKED.from_pretrained()
with torch.no_grad():
_, out = fn.extract_dense_map(img_outdoor)
out = K.color.grayscale_to_rgb(out)
img_in = K.color.rgb_to_bgr(img_outdoor)
elif fn_name == "XFeat":
fn = K.feature.XFeat.from_pretrained()
with torch.no_grad():
preprocessed, _, _ = fn._preprocess_tensor(img_outdoor)
_, K1, _ = fn.net(preprocessed)
out = K.feature.XFeat._get_kpts_heatmap(K1)
# upsample heatmap from preprocessed size back to original image size
out = torch.nn.functional.interpolate(out, img_outdoor.shape[-2:], mode="bilinear", align_corners=False)
out = K.color.grayscale_to_rgb(out)
img_in = K.color.rgb_to_bgr(img_outdoor)
else:
fn = getattr(mod, fn_name)
out = fn(img_in)
out = out - out.min()
out = out / (1e-8 + out.max())
# save the output image
out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
out_np = K.image.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 response function {fn_name}")
if __name__ == "__main__":
main()