We implement an advanced, end-to-end Kornia tutorial and demonstrate how modern, differentiable computer vision can be built entirely in PyTorch. We start by constructing GPU-accelerated, synchronized augmentation pipelines for images, masks, and keypoints, then move into differentiable geometry by optimizing a homography directly through gradient descent. We also show how learned feature matching with LoFTR integrates with Kornia’s RANSAC to estimate robust homographies and produce a simple stitched output, even under constrained or offline-safe conditions. Finally, we ground these ideas in practice by training a lightweight CNN on CIFAR-10 using Kornia’s GPU augmentations, highlighting how research-grade vision pipelines translate naturally into learning systems. Check out the FULL CODES here.
import os, math, time, random, urllib.request from dataclasses import dataclass from typing import Tuple import sys, subprocess def pip_install(pkgs): subprocess.check_call([sys.executable, "-m", "pip", "install", "-q"] + pkgs) pip_install([ "kornia==0.8.2", "torch", "torchvision", "matplotlib", "numpy", "opencv-python-headless" ]) import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torchvision import torchvision.transforms.functional as TF import matplotlib.pyplot as plt import cv2 import kornia import kornia.augmentation as K import kornia.geometry.transform as KG from kornia.geometry.ransac import RANSAC from kornia.feature import LoFTR torch.manual_seed(0) np.random.seed(0) random.seed(0) print("Torch:", torch.__version__) print("Kornia:", kornia.__version__) print("Device:", device)We begin by setting up a fully reproducible environment, installing Kornia and its core dependencies to ensure GPU-accelerated, differentiable computer vision runs smoothly in Google Colab. We then import and organize PyTorch, Kornia, and supporting libraries, establishing a clean foundation for geometry, augmentation, and feature-matching workflows. We set the random seed and select the available compute device so that all subsequent experiments remain deterministic, debuggable, and performance-aware. Check out the FULL CODES here.
def to_tensor_img_uint8(img_bgr_uint8: np.ndarray) -> torch.Tensor: img_rgb = cv2.cvtColor(img_bgr_uint8, cv2.COLOR_BGR2RGB) t = torch.from_numpy(img_rgb).permute(2, 0, 1).float() / 255.0 return t.unsqueeze(0) def show(img_t: torch.Tensor, title: str = "", max_size: int = 900): x = img_t.detach().float().cpu().clamp(0, 1) if x.shape[1] == 1: x = x.repeat(1, 3, 1, 1) x = x[0].permute(1, 2, 0).numpy() h, w = x.shape[:2] scale = min(1.0, max_size / max(h, w)) if scale < 1.0: x = cv2.resize(x, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA) plt.figure(figsize=(7, 5)) plt.imshow(x) plt.axis("off") plt.title(title) plt.show() def show_mask(mask_t: torch.Tensor, title: str = ""): x = mask_t.detach().float().cpu().clamp(0, 1)[0, 0].numpy() plt.figure(figsize=(6, 4)) plt.imshow(x) plt.axis("off") plt.title(title) plt.show() def download(url: str, path: str): os.makedirs(os.path.dirname(path), exist_ok=True) if not os.path.exists(path): urllib.request.urlretrieve(url, path) def safe_download(url: str, path: str) -> bool: try: os.makedirs(os.path.dirname(path), exist_ok=True) if not os.path.exists(path): urllib.request.urlretrieve(url, path) return True except Exception as e: print("Download failed:", e) return False def make_grid_mask(h: int, w: int, cell: int = 32) -> torch.Tensor: yy, xx = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") m = (((yy // cell) % 2) ^ ((xx // cell) % 2)).float() return m.unsqueeze(0).unsqueeze(0) def draw_matches(img0_rgb: np.ndarray, img1_rgb: np.ndarray, pts0: np.ndarray, pts1: np.ndarray, max_draw: int = 200) -> np.ndarray: h0, w0 = img0_rgb.shape[:2] h1, w1 = img1_rgb.shape[:2] out = np.zeros((max(h0, h1), w0 + w1, 3), dtype=np.uint8) out[:h0, :w0] = img0_rgb out[:h1, w0:w0+w1] = img1_rgb n = min(len(pts0), len(pts1), max_draw) if n == 0: return out idx = np.random.choice(len(pts0), size=n, replace=False) if len(pts0) > n else np.arange(n) for i in idx: x0, y0 = pts0[i] x1, y1 = pts1[i] x1_shift = x1 + w0 p0 = (int(round(x0)), int(round(y0))) p1 = (int(round(x1_shift)), int(round(y1))) cv2.circle(out, p0, 2, (255, 255, 255), -1, lineType=cv2.LINE_AA) cv2.circle(out, p1, 2, (255, 255, 255), -1, lineType=cv2.LINE_AA) cv2.line(out, p0, p1, (255, 255, 255), 1, lineType=cv2.LINE_AA) return out def normalize_img_for_loftr(img_rgb01: torch.Tensor) -> torch.Tensor: if img_rgb01.shape[1] == 3: return kornia.color.rgb_to_grayscale(img_rgb01) return img_rgb01We define a set of reusable helper utilities for image conversion, visualization, safe data downloading, and synthetic mask generation, keeping the vision pipeline clean and modular. We also implement robust visualization and matching helpers that allow us to inspect augmented images, masks, and LoFTR correspondences directly during experimentation. We normalize image inputs to the exact tensor formats expected by Kornia and LoFTR, ensuring that all downstream geometry and feature-matching components operate consistently and correctly. Check out the FULL CODES here.
print("n[1] Differentiable augmentations: image + mask + keypoints") B, C, H, W = 1, 3, 256, 384 img = torch.rand(B, C, H, W, device=device) mask = make_grid_mask(H, W, cell=24).to(device) kps = torch.tensor([[ [40.0, 40.0], [W - 50.0, 50.0], [W * 0.6, H * 0.8], [W * 0.25, H * 0.65], ]], device=device) aug = K.AugmentationSequential( K.RandomResizedCrop((224, 224), scale=(0.6, 1.0), ratio=(0.8, 1.25), p=1.0), K.RandomHorizontalFlip(p=0.5), K.RandomRotation(degrees=18.0, p=0.7), K.ColorJiggle(0.2, 0.2, 0.2, 0.1, p=0.8), data_keys=["input", "mask", "keypoints"], same_on_batch=True ).to(device) img_aug, mask_aug, kps_aug = aug(img, mask, kps) print("image:", tuple(img.shape), "->", tuple(img_aug.shape)) print("mask :", tuple(mask.shape), "->", tuple(mask_aug.shape)) print("kps :", tuple(kps.shape), "->", tuple(kps_aug.shape)) print("Example keypoints (before -> after):") print(torch.cat([kps[0], kps_aug[0]], dim=1)) show(img, "Original (synthetic)") show_mask(mask, "Original mask (synthetic)") show(img_aug, "Augmented (synced)") show_mask(mask_aug, "Augmented mask (synced)")We construct a synchronized, fully differentiable augmentation pipeline that applies the same geometric transformations to images, masks, and keypoints on the GPU. We generate synthetic data to clearly demonstrate how spatial consistency is preserved across modalities while still introducing realistic variability through cropping, rotation, flipping, and color jitter. We visualize the before-and-after results to verify that the augmented images, segmentation masks, and keypoints remain perfectly aligned after transformation. Check out the FULL CODES here.
print("n[2] Differentiable homography alignment by optimization") base = torch.rand(1, 1, 240, 320, device=device) show(base, "Base image (grayscale)") true_H_px = torch.eye(3, device=device).unsqueeze(0) true_H_px[:, 0, 2] = 18.0 true_H_px[:, 1, 2] = -12.0 true_H_px[:, 0, 1] = 0.03 true_H_px[:, 1, 0] = -0.02 true_H_px[:, 2, 0] = 1e-4 true_H_px[:, 2, 1] = -8e-5 target = KG.warp_perspective(base, true_H_px, dsize=(base.shape[-2], base.shape[-1]), align_corners=True) show(target, "Target (base warped by true homography)") p = torch.zeros(1, 8, device=device, requires_grad=True) def params_to_H(p8: torch.Tensor) -> torch.Tensor: Bp = p8.shape[0] Hm = torch.eye(3, device=p8.device).unsqueeze(0).repeat(Bp, 1, 1) Hm[:, 0, 0] = 1.0 + p8[:, 0] Hm[:, 0, 1] = p8[:, 1] Hm[:, 0, 2] = p8[:, 2] Hm[:, 1, 0] = p8[:, 3] Hm[:, 1, 1] = 1.0 + p8[:, 4] Hm[:, 1, 2] = p8[:, 5] Hm[:, 2, 0] = p8[:, 6] Hm[:, 2, 1] = p8[:, 7] return Hm opt = torch.optim.Adam([p], lr=0.08) losses = [] for step in range(120): opt.zero_grad(set_to_none=True) H_est = params_to_H(p) pred = KG.warp_perspective(base, H_est, dsize=(base.shape[-2], base.shape[-1]), align_corners=True) loss_photo = (pred - target).abs().mean() loss_reg = 1e-3 * (p ** 2).mean() loss = loss_photo + loss_reg loss.backward() opt.step() losses.append(loss.item()) print("Final loss:", losses[-1]) plt.figure(figsize=(6,4)) plt.plot(losses) plt.title("Homography optimization loss") plt.xlabel("step") plt.ylabel("loss") plt.show() H_est_final = params_to_H(p.detach()) pred_final = KG.warp_perspective(base, H_est_final, dsize=(base.shape[-2], base.shape[-1]), align_corners=True) show(pred_final, "Recovered warp (optimized)") show((pred_final - target).abs(), "Abs error (recovered vs target)") print("True H (pixel):n", true_H_px.squeeze(0).detach().cpu().numpy()) print("Est H:n", H_est_final.squeeze(0).detach().cpu().numpy())We demonstrate that geometric alignment can be treated as a differentiable optimization problem by directly recovering a homography via gradient descent. We first generate a target image by warping a base image with a known homography and then learn the transformation parameters by minimizing a photometric reconstruction loss with regularization. Also, we visualize the optimized warp and error map to confirm that the estimated homography closely matches the ground-truth transformation. Check out the FULL CODES here.
print("n[3] LoFTR matching + RANSAC homography + stitching (403-safe)") data_dir = "https://www.marktechpost.com/content/kornia_demo" os.makedirs(data_dir, exist_ok=True) img0_path = os.path.join(data_dir, "img0.png") img1_path = os.path.join(data_dir, "img1.png") ok0 = safe_download( "https://raw.githubusercontent.com/opencv/opencv/master/samples/data/graf1.png", img0_path ) ok1 = safe_download( "https://raw.githubusercontent.com/opencv/opencv/master/samples/data/graf3.png", img1_path ) if not (ok0 and ok1): print("⚠️ Using synthetic fallback images (no network / blocked downloads)") base_rgb = torch.rand(1, 3, 480, 640, device=device) H_syn = torch.tensor([[ [1.0, 0.05, 40.0], [-0.03, 1.0, 25.0], [1e-4, -8e-5, 1.0] ]], device=device) t0 = base_rgb t1 = KG.warp_perspective(base_rgb, H_syn, dsize=(480, 640), align_corners=True) img0_rgb = (t0[0].permute(1,2,0).detach().cpu().numpy() * 255).astype(np.uint8) img1_rgb = (t1[0].permute(1,2,0).detach().cpu().numpy() * 255).astype(np.uint8) else: img0_bgr = cv2.imread(img0_path, cv2.IMREAD_COLOR) img1_bgr = cv2.imread(img1_path, cv2.IMREAD_COLOR) if img0_bgr is None or img1_bgr is None: raise RuntimeError("Failed to load downloaded images.") img0_rgb = cv2.cvtColor(img0_bgr, cv2.COLOR_BGR2RGB) img1_rgb = cv2.cvtColor(img1_bgr, cv2.COLOR_BGR2RGB) t0 = to_tensor_img_uint8(img0_bgr).to(device) t1 = to_tensor_img_uint8(img1_bgr).to(device) show(t0, "Image 0") show(t1, "Image 1") g0 = normalize_img_for_loftr(t0) g1 = normalize_img_for_loftr(t1) loftr = LoFTR(pretrained="outdoor").to(device).eval() with torch.inference_mode(): correspondences = loftr({"image0": g0, "image1": g1}) mkpts0 = correspondences["keypoints0"] mkpts1 = correspondences["keypoints1"] mconf = correspondences.get("confidence", None) print("Raw matches:", mkpts0.shape[0]) if mkpts0.shape[0] < 8: raise RuntimeError("Too few matches to estimate homography.") if mconf is not None: mconf = mconf.detach() topk = min(2000, mkpts0.shape[0]) idx = torch.topk(mconf, k=topk, largest=True).indices mkpts0 = mkpts0[idx] mkpts1 = mkpts1[idx] print("Kept top matches:", mkpts0.shape[0]) ransac = RANSAC( model_type="homography", inl_th=3.0, batch_size=4096, max_iter=10, confidence=0.999, max_lo_iters=5 ).to(device) with torch.inference_mode(): H01, inliers = ransac(mkpts0, mkpts1) print("Estimated H shape:", tuple(H01.shape)) print("Inliers:", int(inliers.sum().item()), "https://www.marktechpost.com/", int(inliers.numel())) vis = draw_matches( img0_rgb, img1_rgb, mkpts0.detach().cpu().numpy(), mkpts1.detach().cpu().numpy(), max_draw=250 ) plt.figure(figsize=(10,5)) plt.imshow(vis) plt.axis("off") plt.title("LoFTR matches (subset)") plt.show() H01 = H01.unsqueeze(0) if H01.ndim == 2 else H01 warped0 = KG.warp_perspective(t0, H01, dsize=(t1.shape[-2], t1.shape[-1]), align_corners=True) stitched = torch.max(warped0, t1) show(warped0, "Image0 warped into Image1 frame (via RANSAC homography)") show(stitched, "Simple stitched blend (max)")We perform learned feature matching using LoFTR to establish dense correspondences between two images, while ensuring robustness through a network-safe fallback mechanism. We then apply Kornia’s RANSAC to estimate a stable homography from these matches and warp one image into the coordinate frame of the other. We visualize the correspondences and produce a simple stitched result to validate the geometric alignment end-to-end. Check out the FULL CODES here.
print("n[4] Mini training loop with Kornia augmentations (fast subset)") cifar = torchvision.datasets.CIFAR10(root="https://www.marktechpost.com/content/data", train=True, download=True) num_samples = 4096 indices = np.random.permutation(len(cifar))[:num_samples] subset = torch.utils.data.Subset(cifar, indices.tolist()) def collate(batch): imgs = [] labels = [] for im, y in batch: imgs.append(TF.to_tensor(im)) labels.append(y) return torch.stack(imgs, 0), torch.tensor(labels) loader = torch.utils.data.DataLoader( subset, batch_size=256, shuffle=True, num_workers=2, pin_memory=True, collate_fn=collate ) aug_train = K.ImageSequential( K.RandomHorizontalFlip(p=0.5), K.RandomAffine(degrees=12.0, translate=(0.08, 0.08), scale=(0.9, 1.1), p=0.7), K.ColorJiggle(0.2, 0.2, 0.2, 0.1, p=0.8), K.RandomGaussianBlur((3, 3), (0.1, 1.5), p=0.3), ).to(device) class TinyCifarNet(nn.Module): def __init__(self, num_classes=10): super().__init__() self.conv1 = nn.Conv2d(3, 48, 3, padding=1) self.conv2 = nn.Conv2d(48, 96, 3, padding=1) self.conv3 = nn.Conv2d(96, 128, 3, padding=1) self.head = nn.Linear(128, num_classes) def forward(self, x): x = F.relu(self.conv1(x)) x = F.max_pool2d(x, 2) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, 2) x = F.relu(self.conv3(x)) x = x.mean(dim=(-2, -1)) return self.head(x) model = TinyCifarNet().to(device) opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-4) model.train() t_start = time.time() running = [] for it, (xb, yb) in enumerate(loader): xb = xb.to(device, non_blocking=True) yb = yb.to(device, non_blocking=True) xb = aug_train(xb) logits = model(xb) loss = F.cross_entropy(logits, yb) opt.zero_grad(set_to_none=True) loss.backward() opt.step() running.append(loss.item()) if (it + 1) % 10 == 0: print(f"iter {it+1:03d}/{len(loader)} | loss {np.mean(running[-10:]):.4f}") if it >= 39: break print("Done in", round(time.time() - t_start, 2), "sec") plt.figure(figsize=(6,4)) plt.plot(running) plt.title("Training loss (quick demo)") plt.xlabel("iteration") plt.ylabel("loss") plt.show() xb0, yb0 = next(iter(loader)) xb0 = xb0[:8].to(device) xbA = aug_train(xb0) def tile8(x): x = x.detach().cpu().clamp(0,1) grid = torchvision.utils.make_grid(x, nrow=4) return grid.permute(1,2,0).numpy() plt.figure(figsize=(10,5)) plt.imshow(tile8(xb0)) plt.axis("off") plt.title("CIFAR batch (original)") plt.show() plt.figure(figsize=(10,5)) plt.imshow(tile8(xbA)) plt.axis("off") plt.title("CIFAR batch (Kornia-augmented on GPU)") plt.show() print("n✅ Tutorial complete.") print("Next ideas:") print("- Feathered stitching (soft masks) instead of max-blend.") print("- Compare LoFTR vs DISK/LightGlue using kornia.feature.") print("- Multi-scale homography optimization + SSIM/Charbonnier losses.")We demonstrate how Kornia’s GPU-based augmentations integrate directly into a standard training loop by applying them on the fly to a subset of the CIFAR-10 dataset. We train a lightweight convolutional network end-to-end, demonstrating that differentiable augmentations incur minimal overhead while improving data diversity. At last, we visualize original versus augmented batches to confirm that the transformations are applied consistently and efficiently during learning.
In conclusion, we demonstrated that Kornia enables a unified vision workflow where data augmentation, geometric reasoning, feature matching, and learning remain differentiable and GPU-friendly within a single framework. By combining LoFTR matching, RANSAC-based homography estimation, and optimization-driven alignment with a practical training loop, we showed how classical vision and deep learning complement each other rather than compete. It serves as a foundation for extending toward production-grade stitching, robust pose estimation, or large-scale training pipelines, and we emphasize that the same patterns we used here scale naturally to more complex, real-world vision systems.
Check out the FULL CODES here. Also, feel free to follow us on Twitter and don’t forget to join our 100k+ ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
