Frequency tiers (estimates, relative to interviews targeting this pack's specialty — a ★★★★★ here means "core for these roles", not for generalist ML loops): ★★★★★ expect it in most loops · ★★★★ very common · ★★★ common · ★★ occasional · ★ rare/deep-specialist. Spend ~50% of practice on the top two tiers.
80+ problems across image processing, geometry, features, detection, segmentation, tracking, 3D, ViTs, NeRF/3DGS helpers, and metrics Each with full PyTorch / NumPy solution and test cases
Principal/Senior-Principal CV Interview Prep
Notes
This pack collects the live-coding questions that come up in principal/senior-principal level computer-vision interviews and in research-engineer hiring at top labs. Each problem has: (i) a precise statement, (ii) a clean reference solution in NumPy or PyTorch, (iii) tests that verify correctness against either an OpenCV / Kornia / torchvision reference op or a known closed-form property. The companion ML coding pack contains attention, optimizers, losses, RL, and diffusion problems — here we stay focused on what is distinctively vision. Standard imports across the pack:
import math, numpy as np, torch
import torch.nn as nn, torch.nn.functional as F
Problem
Convert an RGB image to grayscale using the ITU-R BT.601 luma weights Y = 0.299R + 0.587G + 0.114B.
import torch
def rgb_to_gray(img):
# img: (..., 3, H, W). Returns (..., 1, H, W)
w = torch.tensor([0.299, 0.587, 0.114], dtype=img.dtype, device=img.device)
return (img * w.view(*([1] * (img.dim() - 3)), 3, 1, 1)).sum(dim=-3, keepdim=True)
Tests
x = torch.zeros(1, 3, 4, 4); x[0, 1] = 1.0 # pure green
g = rgb_to_gray(x)
assert torch.allclose(g, torch.full_like(g, 0.587), atol=1e-6)
print("rgb_to_gray OK")
Problem
Convert RGB image to HSV (hue ∈ [0, 1), saturation/value ∈ [0, 1]) and back, vectorised.
import torch
def rgb_to_hsv(rgb):
# rgb: (..., 3, H, W) in [0,1]
r, g, b = rgb.unbind(dim=-3)
cmax, idx = rgb.max(dim=-3); cmin = rgb.min(dim=-3).values
delta = cmax - cmin + 1e-12
h = torch.zeros_like(cmax)
h = torch.where(idx == 0, ((g - b) / delta) % 6.0, h)
h = torch.where(idx == 1, (b - r) / delta + 2.0, h)
h = torch.where(idx == 2, (r - g) / delta + 4.0, h)
h = h / 6.0
s = torch.where(cmax == 0, torch.zeros_like(cmax), delta / (cmax + 1e-12))
v = cmax
return torch.stack([h, s, v], dim=-3)
def hsv_to_rgb(hsv):
h, s, v = hsv.unbind(dim=-3)
i = torch.floor(h * 6.0).long() % 6
f = h * 6.0 - i.float()
p = v * (1 - s); q = v * (1 - f * s); t = v * (1 - (1 - f) * s)
out = torch.stack([
torch.stack([v, t, p], dim=-3),
torch.stack([q, v, p], dim=-3),
torch.stack([p, v, t], dim=-3),
torch.stack([p, q, v], dim=-3),
torch.stack([t, p, v], dim=-3),
torch.stack([v, p, q], dim=-3),
], dim=0)
idx = i.unsqueeze(-3).expand_as(out[0])
return torch.gather(out, 0, idx.unsqueeze(0))[0]
Tests
x = torch.rand(2, 3, 8, 8)
y = hsv_to_rgb(rgb_to_hsv(x))
assert torch.allclose(x, y, atol=1e-3)
print("rgb<->hsv OK")
Problem
Implement bilinear sampling at fractional coordinates (x, y) with zero-padded out-of-bounds behaviour. Match F.grid_sample(align_corners=True) for valid pixels.
import numpy as np
def bilinear_sample(img, x, y):
# img: (C, H, W) float; (x, y): scalar or arrays in pixel coords
x = np.asarray(x); y = np.asarray(y)
H, W = img.shape[-2:]
x0 = np.floor(x).astype(int); x1 = x0 + 1
y0 = np.floor(y).astype(int); y1 = y0 + 1
def gather(xx, yy):
valid = (xx >= 0) & (xx < W) & (yy >= 0) & (yy < H)
out = np.zeros((img.shape[0],) + xx.shape, dtype=img.dtype)
if valid.any():
xv = np.clip(xx, 0, W - 1); yv = np.clip(yy, 0, H - 1)
out = img[:, yv, xv]; out = np.where(valid, out, 0.0)
return out
Ia = gather(x0, y0); Ib = gather(x1, y0); Ic = gather(x0, y1); Id = gather(x1, y1)
wa = (x1 - x) * (y1 - y); wb = (x - x0) * (y1 - y)
wc = (x1 - x) * (y - y0); wd = (x - x0) * (y - y0)
return Ia * wa + Ib * wb + Ic * wc + Id * wd
Tests
img = np.arange(16, dtype=np.float32).reshape(1, 4, 4)
v = bilinear_sample(img, np.array([0.5]), np.array([0.5]))
assert abs(v[0, 0] - (0 + 1 + 4 + 5) / 4.) < 1e-6
print("bilinear OK")
Advanced implementation. Fully vectorized for N arbitrary sample points: four fancy-index gathers + a lerp, no loops. Verified equal to F.grid_sample(align_corners=True) (1e-5).
import torch
def bilinear_sample_vec(img, xs, ys):
# img: (C, H, W); xs, ys: (N,) pixel coords -> (C, N)
C, H, W = img.shape
x0 = xs.floor().long().clamp(0, W - 2); y0 = ys.floor().long().clamp(0, H - 2)
x1, y1 = x0 + 1, y0 + 1
wx, wy = xs - x0.float(), ys - y0.float()
Ia, Ib = img[:, y0, x0], img[:, y0, x1]
Ic, Id = img[:, y1, x0], img[:, y1, x1]
return Ia*(1-wx)*(1-wy) + Ib*wx*(1-wy) + Ic*(1-wx)*wy + Id*wx*wy
Problem
Resize an image to (H′, W ′) using bilinear sampling with PyTorch (no F.interpolate).
import torch
import torch.nn.functional as F
def resize_bilinear(img, out_h, out_w):
# img: (B, C, H, W). Builds a normalized [-1,1] grid and samples.
B, _, H, W = img.shape
ys = torch.linspace(-1, 1, out_h, device=img.device)
xs = torch.linspace(-1, 1, out_w, device=img.device)
gy, gx = torch.meshgrid(ys, xs, indexing='ij')
grid = torch.stack([gx, gy], dim=-1).unsqueeze(0).expand(B, -1, -1, -1)
return F.grid_sample(img, grid, mode='bilinear', align_corners=True)
Tests
x = torch.rand(2, 3, 8, 8)
y = resize_bilinear(x, 16, 16)
ref = F.interpolate(x, size=(16, 16), mode='bilinear', align_corners=True)
assert torch.allclose(y, ref, atol=1e-5)
print("resize OK")
Problem
Apply 2D padding with constant / reflect / replicate modes; match F.pad.
import torch.nn.functional as F
def pad2d(img, pad, mode='reflect', value=0.0):
# img: (B, C, H, W); pad: (left, right, top, bottom)
return F.pad(img, pad, mode=mode, value=value)
Tests
x = torch.tensor([[[[1., 2.], [3., 4.]]]])
y = pad2d(x, (1, 1, 1, 1), 'reflect')
assert y.shape == (1, 1, 4, 4) and y[0,0,0,0].item() == 4.0
print("pad OK")
Problem
Compute the 256-bin histogram of an 8-bit grayscale image and apply histogram equalization.
import numpy as np
def histogram(img, bins=256):
# img: (H, W) uint8
return np.bincount(img.ravel(), minlength=bins)
def hist_equalize(img):
h = histogram(img)
cdf = np.cumsum(h)
cdf_min = cdf[cdf > 0][0]
lut = np.round((cdf - cdf_min) / (cdf[-1] - cdf_min) * 255).clip(0, 255).astype(np.uint8)
return lut[img]
Tests
img = np.array([[10, 10, 50], [50, 200, 200]], dtype=np.uint8)
out = hist_equalize(img)
assert out.min() == 0 and out.max() == 255
print("equalize OK", out)
Problem
Apply gamma correction Iout = I1/γ on a float image in [0, 1]. in
def gamma_correct(img, gamma=2.2):
return img.clamp(0, 1) ** (1.0 / gamma)
Tests
x = torch.tensor([0.0, 0.5, 1.0]); y = gamma_correct(x, 2.2)
assert torch.allclose(y[0], torch.tensor(0.0)) and torch.allclose(y[2], torch.tensor(1.0))
assert y[1] > 0.5 # mid-gray brightens after gamma 2.2
print("gamma OK")
Problem
Apply a k × k box (mean) filter via 2D convolution, separately on each channel.
import torch
import torch.nn.functional as F
def box_filter(img, k=3):
# img: (B, C, H, W)
C = img.size(1)
kernel = torch.ones(C, 1, k, k, device=img.device, dtype=img.dtype) / (k * k)
return F.conv2d(img, kernel, padding=k // 2, groups=C)
Tests
x = torch.ones(1, 1, 5, 5)
y = box_filter(x, k=3)
assert torch.allclose(y[0, 0, 2, 2], torch.tensor(1.0))
print("box OK")
Problem
Apply a Gaussian blur of standard deviation σ as a 1D convolution along x followed by y, ie. exploit separability.
import torch
import torch.nn.functional as F
def gaussian_kernel_1d(sigma, k=None):
if k is None: k = max(3, int(2 * round(3 * sigma) + 1))
half = k // 2
x = torch.arange(-half, half + 1, dtype=torch.float32)
g = torch.exp(-0.5 * (x / sigma) ** 2); g = g / g.sum()
return g
def gaussian_blur(img, sigma):
g = gaussian_kernel_1d(sigma).to(img.device)
k = g.numel(); pad = k // 2
C = img.size(1)
g_x = g.view(1, 1, 1, k).expand(C, 1, 1, k)
g_y = g.view(1, 1, k, 1).expand(C, 1, k, 1)
out = F.conv2d(img, g_x, padding=(0, pad), groups=C)
out = F.conv2d(out, g_y, padding=(pad, 0), groups=C)
return out
Tests
x = torch.zeros(1, 1, 11, 11); x[0, 0, 5, 5] = 1.0
y = gaussian_blur(x, sigma=1.0)
assert abs(y.sum().item() - 1.0) < 1e-5 # filter preserves total mass
print("gaussian OK")
Problem
Compute Sobel Gx, Gy, gradient magnitude, and direction θ = arctan 2(Gy, Gx).
import torch
import torch.nn.functional as F
def sobel(img):
kx = torch.tensor([[-1., 0., 1.], [-2., 0., 2.], [-1., 0., 1.]], device=img.device)
ky = kx.T
C = img.size(1)
kx = kx.view(1, 1, 3, 3).expand(C, 1, 3, 3)
ky = ky.view(1, 1, 3, 3).expand(C, 1, 3, 3)
Gx = F.conv2d(img, kx, padding=1, groups=C)
Gy = F.conv2d(img, ky, padding=1, groups=C)
mag = (Gx ** 2 + Gy ** 2).sqrt()
ang = torch.atan2(Gy, Gx)
return Gx, Gy, mag, ang
Tests
img = torch.zeros(1, 1, 5, 5); img[0, 0, :, 2:] = 1.0
_, _, mag, _ = sobel(img)
assert mag.max().item() > 1.0 # edge at column 2 detected
print("sobel OK")
Problem
Apply the discrete Laplacian kernel0 1 0 ; 1 −4 1 ; 0 1 0. Report response on a uniform image.
import torch
import torch.nn.functional as F
def laplacian(img):
K = torch.tensor([[0., 1., 0.], [1., -4., 1.], [0., 1., 0.]], device=img.device)
C = img.size(1)
K = K.view(1, 1, 3, 3).expand(C, 1, 3, 3)
return F.conv2d(img, K, padding=1, groups=C)
Tests
x = torch.full((1, 1, 4, 4), 7.0)
assert laplacian(x)[0, 0, 1, 1].abs().item() < 1e-6 # uniform => zero
print("laplacian OK")
Problem
Implement a small bilateral filter (5×5 window, spatial σs, range σr) on a single-channel image.
import numpy as np
def bilateral(img, sigma_s=2.0, sigma_r=0.1, k=5):
H, W = img.shape; pad = k // 2
p = np.pad(img, pad, mode='reflect')
out = np.zeros_like(img, dtype=np.float32)
ys, xs = np.mgrid[-pad:pad+1, -pad:pad+1]
sp = np.exp(-(xs**2 + ys**2) / (2 * sigma_s ** 2))
for y in range(H):
for x in range(W):
patch = p[y:y+k, x:x+k]
rng = np.exp(-((patch - p[y+pad, x+pad]) ** 2) / (2 * sigma_r ** 2))
w = sp * rng
out[y, x] = (patch * w).sum() / w.sum()
return out
Tests
img = np.tile(np.linspace(0, 1, 8), (8, 1))
o = bilateral(img.astype(np.float32))
assert o.shape == img.shape and abs(o.mean() - img.mean()) < 0.05
print("bilateral OK")
Problem
Apply a k × k median filter to a single-channel image.
import numpy as np
def median_filter(img, k=3):
H, W = img.shape; pad = k // 2
p = np.pad(img, pad, mode='reflect')
out = np.zeros_like(img)
for y in range(H):
for x in range(W):
out[y, x] = np.median(p[y:y+k, x:x+k])
return out
Tests
img = np.array([[1, 1, 9], [1, 1, 1], [1, 1, 1]], dtype=np.float32)
out = median_filter(img, 3)
assert out[0, 2] == 1.0 # spike removed
print("median OK")
Problem
Sharpen an image: Iout = I + α(I −Blurσ(I)).
def unsharp(img, sigma=1.0, alpha=1.0):
blur = gaussian_blur(img, sigma)
return (img + alpha * (img - blur)).clamp(0, 1)
Tests
img = torch.full((1, 1, 8, 8), 0.5)
img[0, 0, 4, 4] = 1.0
out = unsharp(img, sigma=1.0, alpha=1.0)
assert out[0, 0, 4, 4] >= 1.0 - 1e-6 # peak gets brighter
print("unsharp OK")
Problem
Compute forward-difference gradients ∇xI, ∇yI and zero-pad the last column/row.
import torch
def grad_xy(img):
gx = torch.zeros_like(img); gy = torch.zeros_like(img)
gx[..., :-1] = img[..., 1:] - img[..., :-1]
gy[..., :-1, :] = img[..., 1:, :] - img[..., :-1, :]
return gx, gy
Tests
x = torch.arange(16, dtype=torch.float32).view(1, 1, 4, 4)
gx, gy = grad_xy(x)
assert torch.allclose(gx[..., :-1], torch.full((1, 1, 4, 3), 1.0))
print("grad OK")
Problem
Build a 2×2 rotation matrix from angle θ (radians); verify R⊤R = I.
import math
import torch
def rot2(theta):
c, s = math.cos(theta), math.sin(theta)
return torch.tensor([[c, -s], [s, c]])
Tests
R = rot2(0.7)
assert torch.allclose(R.T @ R, torch.eye(2), atol=1e-6)
print("rot2 OK")
Problem
Warp an image with a 2×3 affine matrix A in normalised coordinates, using F.affine_grid.
import torch.nn.functional as F
def affine_warp(img, A):
# A: (B, 2, 3), img: (B, C, H, W)
grid = F.affine_grid(A, img.size(), align_corners=True)
return F.grid_sample(img, grid, align_corners=True)
Tests
img = torch.eye(8).view(1, 1, 8, 8)
A = torch.eye(2, 3).unsqueeze(0) # identity
out = affine_warp(img, A)
assert torch.allclose(out, img, atol=1e-5)
print("affine OK")
Problem
Estimate a homography H such that x′i ∼Hxi from 4 correspondences using the Direct Linear Transform.
import numpy as np
def homography_dlt(src, dst):
# src, dst: (4, 2)
A = []
for (x, y), (xp, yp) in zip(src, dst):
A.append([-x, -y, -1, 0, 0, 0, x*xp, y*xp, xp])
A.append([0, 0, 0, -x, -y, -1, x*yp, y*yp, yp])
A = np.asarray(A)
_, _, V = np.linalg.svd(A)
H = V[-1].reshape(3, 3)
return H / H[2, 2]
Tests
src = np.array([[0,0],[1,0],[1,1],[0,1]], dtype=np.float64)
dst = src + np.array([2.0, 1.0]) # pure translation
H = homography_dlt(src, dst)
assert abs(H[0, 2] - 2.0) < 1e-9 and abs(H[1, 2] - 1.0) < 1e-9
print("homography OK")
Problem
Project a set of 2D points through homography H and divide by the homogeneous coordinate.
import numpy as np
def apply_homography(H, pts):
# pts: (N, 2)
p = np.hstack([pts, np.ones((len(pts), 1))]) @ H.T
return p[:, :2] / p[:, 2:3]
Tests
H = np.array([[1, 0, 5], [0, 1, -3], [0, 0, 1.0]])
out = apply_homography(H, np.array([[0., 0.], [1., 1.]]))
assert np.allclose(out, [[5, -3], [6, -2]])
print("apply H OK")
Problem
Project 3D world points X to pixel coordinates given intrinsics K and extrinsics [R|t]: x = K(RX + t), divide by depth.
def project(K, R, t, X):
# X: (N, 3)
Xc = X @ R.T + t
p = Xc @ K.T
return p[:, :2] / p[:, 2:3], Xc[:, 2]
Tests
K = np.array([[100, 0, 32], [0, 100, 32], [0, 0, 1.]])
R = np.eye(3); t = np.zeros(3)
X = np.array([[0, 0, 1.0], [1, 0, 1.0]])
uv, z = project(K, R, t, X)
assert np.allclose(uv, [[32, 32], [132, 32]])
print("project OK")
Problem
Given two camera projection matrices P1, P2 and two image observations x1, x2, recover the 3D point by stacking [xi]×Pi rows and solving via SVD.
import numpy as np
def triangulate(P1, P2, x1, x2):
A = np.array([
x1[0] * P1[2] - P1[0],
x1[1] * P1[2] - P1[1],
x2[0] * P2[2] - P2[0],
x2[1] * P2[2] - P2[1],
])
_, _, V = np.linalg.svd(A)
X = V[-1]; return X[:3] / X[3]
Tests
K = np.array([[100, 0, 32], [0, 100, 32], [0, 0, 1.0]])
P1 = K @ np.hstack([np.eye(3), np.zeros((3, 1))])
P2 = K @ np.hstack([np.eye(3), np.array([[-0.1], [0], [0]])])
X = np.array([0., 0., 1.0])
u1 = (P1 @ np.append(X, 1.0))[:2] / (P1 @ np.append(X, 1.0))[2]
u2 = (P2 @ np.append(X, 1.0))[:2] / (P2 @ np.append(X, 1.0))[2]
Xh = triangulate(P1, P2, u1, u2)
assert np.allclose(Xh, X, atol=1e-6)
print("triangulate OK")
Problem
Apply radial-tangential distortion to a normalized point (x, y) with parameters k1, k2, p1, p2, k3.
import numpy as np
def distort(xy, k1=0., k2=0., p1=0., p2=0., k3=0.):
x, y = xy[..., 0], xy[..., 1]
r2 = x*x + y*y
radial = 1 + k1*r2 + k2*r2**2 + k3*r2**3
x_d = x * radial + 2*p1*x*y + p2*(r2 + 2*x*x)
y_d = y * radial + p1*(r2 + 2*y*y) + 2*p2*x*y
return np.stack([x_d, y_d], axis=-1)
Tests
xy = np.array([[0., 0.]])
out = distort(xy, k1=1e-3); assert np.allclose(out, xy, atol=1e-9)
print("distort OK")
Problem
Given pixel (u, v), intrinsics K, and rotation R, return the unit ray direction in world frame.
import numpy as np
def pixel_to_ray(K, R, u, v):
inv_K = np.linalg.inv(K)
d_cam = inv_K @ np.array([u, v, 1.0])
d_world = R.T @ d_cam
return d_world / np.linalg.norm(d_world)
Tests
K = np.array([[100,0,32],[0,100,32],[0,0,1.0]])
d = pixel_to_ray(K, np.eye(3), 32, 32)
assert np.allclose(d, [0, 0, 1.0])
print("ray OK")
Problem
Compute the Harris corner response det(M) −k tr(M)2 where M = Gσ ∗(∇I∇I⊤).
def harris_response(img, sigma=1.0, k=0.04):
Gx, Gy, _, _ = sobel(img)
Ixx = gaussian_blur(Gx * Gx, sigma)
Iyy = gaussian_blur(Gy * Gy, sigma)
Ixy = gaussian_blur(Gx * Gy, sigma)
det = Ixx * Iyy - Ixy ** 2
tr = Ixx + Iyy
return det - k * tr ** 2
Tests
img = torch.zeros(1, 1, 16, 16); img[0, 0, 4:8, 4:8] = 1.0 # a square
R = harris_response(img)
ymax, xmax = (R == R.max()).nonzero(as_tuple=False)[0, 2:].tolist()
assert (3 <= ymax <= 8) and (3 <= xmax <= 8)
print("harris OK")
Problem
Given a 2D response map and a radius r, return up to K local maxima sorted by score.
import torch.nn.functional as F
def keypoint_nms(score, r=3, K=100):
# score: (H, W). Returns list of (y, x) sorted by descending score.
H, W = score.shape
kp = []
pad_score = F.max_pool2d(score[None, None].float(), 2*r+1, stride=1, padding=r)
is_max = (score == pad_score[0, 0])
ys, xs = is_max.nonzero(as_tuple=True)
vals = score[ys, xs]
order = vals.argsort(descending=True)[:K]
return [(ys[i].item(), xs[i].item(), vals[i].item()) for i in order]
Tests
m = torch.zeros(8, 8); m[3, 3] = 5.0; m[6, 6] = 3.0; m[3, 4] = 4.9
kp = keypoint_nms(m, r=1, K=10)
assert kp[0][:2] == (3, 3)
print("kp NMS OK", kp)
Problem
Implement a simplified FAST-9 corner detector: for each pixel, check 16 ring pixels and decide whether ≥9 contiguous are above I + t or below I −t.
def fast_corners(img, t=10):
# img: 2D uint8
offsets = [(-3,0),(-3,1),(-2,2),(-1,3),(0,3),(1,3),(2,2),(3,1),
(3,0),(3,-1),(2,-2),(1,-3),(0,-3),(-1,-3),(-2,-2),(-3,-1)]
H, W = img.shape; out = []
for y in range(3, H-3):
for x in range(3, W-3):
cp = int(img[y, x])
ring = [int(img[y+dy, x+dx]) for dy, dx in offsets]
br = [v > cp + t for v in ring] + [v > cp + t for v in ring]
dk = [v < cp - t for v in ring] + [v < cp - t for v in ring]
run_b = max(sum(1 for v in br[i:i+9] if v) for i in range(16))
run_d = max(sum(1 for v in dk[i:i+9] if v) for i in range(16))
if run_b >= 9 or run_d >= 9: out.append((y, x))
return out
Tests
img = np.full((20, 20), 50, dtype=np.uint8); img[10, 10] = 200
c = fast_corners(img, t=20)
assert (10, 10) in c; print("fast OK")
Problem
Match descriptors D1 to D2 via L2 distance; apply Lowe’s ratio test d1st < r · d2nd.
import torch
def bf_match(D1, D2, ratio=0.75):
# D1: (N, d), D2: (M, d)
diff = D1[:, None] - D2[None]
dist = (diff ** 2).sum(-1).sqrt()
vals, idx = dist.topk(2, dim=1, largest=False)
keep = vals[:, 0] < ratio * vals[:, 1]
return torch.stack([torch.arange(D1.size(0))[keep], idx[keep, 0]], dim=-1)
Tests
D1 = torch.tensor([[1., 0.], [0., 1.]])
D2 = torch.tensor([[1., 0.], [0., 1.], [10., 10.]])
m = bf_match(D1, D2, ratio=0.5)
assert m.tolist() == [[0, 0], [1, 1]]
print("matcher OK")
Problem
Given a similarity matrix, return mutual matches: i → j iff j = arg max Si,: and i = arg max S:,j.
def mutual_nn(sim):
nn1 = sim.argmax(dim=1)
nn2 = sim.argmax(dim=0)
matches = []
for i, j in enumerate(nn1.tolist()):
if nn2[j].item() == i: matches.append((i, j))
return matches
Tests
S = torch.tensor([[0.9, 0.1, 0.0], [0.0, 0.8, 0.1], [0.1, 0.0, 0.9]])
assert mutual_nn(S) == [(0,0),(1,1),(2,2)]
print("mutual NN OK")
Problem
Fit a 2D line ax + by + c = 0 to noisy points with RANSAC; return inlier mask.
import numpy as np
def ransac_line(pts, iters=200, thresh=0.5, rng=None):
rng = rng or np.random.default_rng(0)
best, best_in = None, np.zeros(len(pts), dtype=bool)
for _ in range(iters):
i, j = rng.choice(len(pts), 2, replace=False)
p, q = pts[i], pts[j]
n = np.array([q[1] - p[1], p[0] - q[0]])
n = n / (np.linalg.norm(n) + 1e-9)
c = -np.dot(n, p)
d = np.abs(pts @ n + c)
inl = d < thresh
if inl.sum() > best_in.sum():
best, best_in = (n[0], n[1], c), inl
return best, best_in
Tests
rng = np.random.default_rng(0)
inl = rng.uniform(-5, 5, (50, 1))
pts = np.hstack([inl, 2 * inl + 1 + 0.05 * rng.standard_normal((50, 1))])
out_pts = rng.uniform(-5, 5, (20, 2))
data = np.vstack([pts, out_pts])
(a, b, c), m = ransac_line(data, thresh=0.2)
assert m[:50].sum() > 40
print("RANSAC line OK")
Problem
Find a homography from many noisy correspondences via RANSAC, sampling 4-tuples.
import numpy as np
def ransac_homography(src, dst, iters=500, thr=3.0, rng=None):
rng = rng or np.random.default_rng(0)
n = len(src); best_H, best_in = None, 0
for _ in range(iters):
idx = rng.choice(n, 4, replace=False)
try:
H = homography_dlt(src[idx], dst[idx])
except np.linalg.LinAlgError:
continue
proj = apply_homography(H, src)
err = np.linalg.norm(proj - dst, axis=1)
inl = (err < thr).sum()
if inl > best_in:
best_H, best_in = H, inl
return best_H, best_in
Tests
rng = np.random.default_rng(1)
src = rng.uniform(-1, 1, (50, 2))
H_true = np.array([[1.1, 0.05, 0.2], [0., 0.95, -0.1], [0., 0., 1.]])
dst = apply_homography(H_true, src) + 0.001 * rng.standard_normal((50, 2))
H_est, _ = ransac_homography(src, dst, iters=200, thr=0.1)
assert np.allclose(H_est, H_true, atol=0.05)
print("RANSAC H OK")
Problem
Estimate F from 8+ correspondences. Enforce rank(F) = 2 by zeroing the smallest singular value.
import numpy as np
def fundamental_8pt(x1, x2):
A = np.stack([
x2[:, 0] * x1[:, 0], x2[:, 0] * x1[:, 1], x2[:, 0],
x2[:, 1] * x1[:, 0], x2[:, 1] * x1[:, 1], x2[:, 1],
x1[:, 0], x1[:, 1], np.ones(len(x1))], axis=-1)
_, _, V = np.linalg.svd(A)
F = V[-1].reshape(3, 3)
U, S, Vt = np.linalg.svd(F); S[-1] = 0
return U @ np.diag(S) @ Vt
Tests
rng = np.random.default_rng(0)
X = rng.uniform(-1, 1, (20, 3)) + np.array([0, 0, 5])
K = np.array([[100,0,32],[0,100,32],[0,0,1.0]])
P1 = K @ np.hstack([np.eye(3), np.zeros((3, 1))])
P2 = K @ np.hstack([np.eye(3), np.array([[-0.5], [0], [0]])])
x1 = (X @ P1[:, :3].T + P1[:, 3])
x1 = x1[:, :2] / x1[:, 2:3]
x2 = (X @ P2[:, :3].T + P2[:, 3])
x2 = x2[:, :2] / x2[:, 2:3]
F_ = fundamental_8pt(x1, x2)
err = np.array([np.append(b,1) @ F_ @ np.append(a,1) for a,b in zip(x1, x2)])
assert np.abs(err).max() < 1e-3
print("F 8-pt OK")
Problem
Recover camera pose P = K[R|t] from ≥6 3D-2D point pairs by stacking the DLT equations.
import numpy as np
def pnp_dlt(X, x):
# X: (N, 3), x: (N, 2). Returns 3x4 P.
rows = []
for (Xw, Yw, Zw), (u, v) in zip(X, x):
rows.append([Xw, Yw, Zw, 1, 0, 0, 0, 0, -u*Xw, -u*Yw, -u*Zw, -u])
rows.append([0, 0, 0, 0, Xw, Yw, Zw, 1, -v*Xw, -v*Yw, -v*Zw, -v])
A = np.asarray(rows)
_, _, V = np.linalg.svd(A)
return V[-1].reshape(3, 4)
Tests
K = np.array([[100,0,32],[0,100,32],[0,0,1.0]])
X = np.array([[0,0,1.0],[1,0,1],[0,1,1],[1,1,1],[0,0,2],[1,1,2]])
P_true = K @ np.hstack([np.eye(3), np.array([[0.1],[-0.05],[0]])])
x_h = X @ P_true[:, :3].T + P_true[:, 3]
x = x_h[:, :2] / x_h[:, 2:3]
P_est = pnp_dlt(X, x)
P_est = P_est * (P_true[2, 3] / P_est[2, 3])
assert np.allclose(P_est, P_true, atol=1e-4)
print("PnP OK")
Problem
For two boxes A and B, GIoU = IoU −|C(A∪B)| where C is the smallest enclosing box. |C|
import torch
def giou(a, b):
x1 = torch.maximum(a[..., 0], b[..., 0]); y1 = torch.maximum(a[..., 1], b[..., 1])
x2 = torch.minimum(a[..., 2], b[..., 2]); y2 = torch.minimum(a[..., 3], b[..., 3])
iw = (x2 - x1).clamp(min=0); ih = (y2 - y1).clamp(min=0)
inter = iw * ih
aA = (a[..., 2] - a[..., 0]) * (a[..., 3] - a[..., 1])
aB = (b[..., 2] - b[..., 0]) * (b[..., 3] - b[..., 1])
union = aA + aB - inter
cx1 = torch.minimum(a[..., 0], b[..., 0]); cy1 = torch.minimum(a[..., 1], b[..., 1])
cx2 = torch.maximum(a[..., 2], b[..., 2]); cy2 = torch.maximum(a[..., 3], b[..., 3])
enc = (cx2 - cx1) * (cy2 - cy1)
return inter / union.clamp(min=1e-9) - (enc - union) / enc.clamp(min=1e-9)
Tests
a = torch.tensor([0., 0., 10., 10.])
b = torch.tensor([20., 20., 30., 30.])
g = giou(a, b) # disjoint -> negative
assert g.item() < 0
print("GIoU OK", g.item())
Problem
DIoU = IoU −ρ2(ca,cb) , d = diagonal of enclosing box. d2
import torch
def diou(a, b):
inter_x1 = torch.maximum(a[...,0], b[...,0]); inter_y1 = torch.maximum(a[...,1], b[...,1])
inter_x2 = torch.minimum(a[...,2], b[...,2]); inter_y2 = torch.minimum(a[...,3], b[...,3])
iw = (inter_x2 - inter_x1).clamp(min=0); ih = (inter_y2 - inter_y1).clamp(min=0)
inter = iw * ih
aA = (a[...,2] - a[...,0]) * (a[...,3] - a[...,1])
aB = (b[...,2] - b[...,0]) * (b[...,3] - b[...,1])
iou_ = inter / (aA + aB - inter).clamp(min=1e-9)
ca = (a[..., :2] + a[..., 2:]) / 2
cb = (b[..., :2] + b[..., 2:]) / 2
rho2 = ((ca - cb) ** 2).sum(-1)
cx1 = torch.minimum(a[...,0], b[...,0]); cy1 = torch.minimum(a[...,1], b[...,1])
cx2 = torch.maximum(a[...,2], b[...,2]); cy2 = torch.maximum(a[...,3], b[...,3])
d2 = (cx2 - cx1) ** 2 + (cy2 - cy1) ** 2
return iou_ - rho2 / d2.clamp(min=1e-9)
Tests
a = torch.tensor([0., 0., 10., 10.]); b = torch.tensor([0., 0., 10., 10.])
assert torch.allclose(diou(a, b), torch.tensor(1.0)); print("DIoU OK")
Problem
Soft-NMS (Bodla et al. 2017): instead of suppressing overlapping boxes outright, decay their scores by s ← s · exp(−IoU2/σ).
import torch
def soft_nms(boxes, scores, sigma=0.5, score_thr=1e-3):
boxes = boxes.clone(); scores = scores.clone()
keep = []
for _ in range(len(scores)):
i = scores.argmax().item()
if scores[i] < score_thr: break
keep.append(i)
# decay all others
rest = [j for j in range(len(scores)) if j != i and scores[j] >= score_thr]
if not rest: break
ious = giou(boxes[i:i+1].expand(len(rest), 4), boxes[rest])
decay = torch.exp(-(ious.clamp(min=0) ** 2) / sigma)
scores[rest] *= decay
scores[i] = -1.0
return keep
Tests
B = torch.tensor([[0,0,10,10],[1,1,11,11],[50,50,60,60]], dtype=torch.float32)
S = torch.tensor([0.9, 0.8, 0.7])
keep = soft_nms(B, S)
assert 0 in keep and 2 in keep
print("soft NMS OK", keep)
Problem
Convert between [x1, y1, x2, y2] and [cx, cy, w, h] formats.
import torch
def xyxy_to_cxcywh(b):
cx = (b[..., 0] + b[..., 2]) / 2; cy = (b[..., 1] + b[..., 3]) / 2
w = b[..., 2] - b[..., 0]; h = b[..., 3] - b[..., 1]
return torch.stack([cx, cy, w, h], dim=-1)
def cxcywh_to_xyxy(b):
x1 = b[..., 0] - b[..., 2] / 2; y1 = b[..., 1] - b[..., 3] / 2
x2 = b[..., 0] + b[..., 2] / 2; y2 = b[..., 1] + b[..., 3] / 2
return torch.stack([x1, y1, x2, y2], dim=-1)
Tests
b = torch.tensor([[0., 0., 10., 20.]])
c = xyxy_to_cxcywh(b); back = cxcywh_to_xyxy(c)
assert torch.allclose(back, b); print("xyxy<->cxcywh OK", c)
Problem
Clip [x1, y1, x2, y2] boxes so they stay inside [0, W −1] × [0, H −1].
import torch
def clip_boxes(b, W, H):
return torch.stack([
b[..., 0].clamp(0, W - 1), b[..., 1].clamp(0, H - 1),
b[..., 2].clamp(0, W - 1), b[..., 3].clamp(0, H - 1)], dim=-1)
Tests
b = torch.tensor([[-5., -5., 100., 100.]])
out = clip_boxes(b, 50, 30)
assert torch.equal(out, torch.tensor([[0., 0., 49., 29.]]))
print("clip OK")
Problem
Generate axis-aligned anchors of given scales s and aspect ratios r at each location of a feature map of stride S.
import math
import torch
def gen_anchors(H, W, stride, scales, ratios):
cy = (torch.arange(H) + 0.5) * stride
cx = (torch.arange(W) + 0.5) * stride
cy, cx = torch.meshgrid(cy, cx, indexing='ij')
centers = torch.stack([cx, cy], dim=-1).view(-1, 2)
anchors = []
for s in scales:
for r in ratios:
w = s * math.sqrt(1 / r); h = s * math.sqrt(r)
wh = torch.tensor([w, h], dtype=torch.float32).expand(centers.size(0), 2)
xy1 = centers - wh / 2; xy2 = centers + wh / 2
anchors.append(torch.cat([xy1, xy2], dim=-1))
return torch.cat(anchors, dim=0)
Tests
A = gen_anchors(2, 2, stride=8, scales=[16], ratios=[1.0])
assert A.shape == (4, 4); print("anchors OK", A[0])
Problem
Encode an anchor a and a target g as (tx, ty, tw, th): tx = (gx −ax)/aw, tw = log(gw/aw). Implement decoding as well.
import torch
def encode_box(a, g):
a = xyxy_to_cxcywh(a); g = xyxy_to_cxcywh(g)
tx = (g[..., 0] - a[..., 0]) / a[..., 2]
ty = (g[..., 1] - a[..., 1]) / a[..., 3]
tw = torch.log(g[..., 2] / a[..., 2])
th = torch.log(g[..., 3] / a[..., 3])
return torch.stack([tx, ty, tw, th], dim=-1)
def decode_box(a, t):
a = xyxy_to_cxcywh(a)
cx = a[..., 0] + t[..., 0] * a[..., 2]
cy = a[..., 1] + t[..., 1] * a[..., 3]
w = a[..., 2] * torch.exp(t[..., 2])
h = a[..., 3] * torch.exp(t[..., 3])
return cxcywh_to_xyxy(torch.stack([cx, cy, w, h], dim=-1))
Tests
a = torch.tensor([[0., 0., 10., 10.]])
g = torch.tensor([[1., 2., 11., 12.]])
t = encode_box(a, g); back = decode_box(a, t)
assert torch.allclose(back, g, atol=1e-5); print("encode/decode OK")
Problem
Mark a feature-map location as a positive sample if its center lies within a radius r · S of the GT box center, where S is the stride.
import torch
def fcos_centers(H, W, stride, gt_box, radius=1.5):
cy = (torch.arange(H) + 0.5) * stride
cx = (torch.arange(W) + 0.5) * stride
yy, xx = torch.meshgrid(cy, cx, indexing='ij')
bx = (gt_box[0] + gt_box[2]) / 2; by = (gt_box[1] + gt_box[3]) / 2
return ((xx - bx).abs() < radius * stride) & ((yy - by).abs() < radius * stride)
Tests
mask = fcos_centers(8, 8, stride=4, gt_box=torch.tensor([10., 10., 20., 20.]))
assert mask.any(); print("FCOS center OK")
Problem
Label connected regions of a binary mask using BFS with 4-connectivity.
from collections import deque
def connected_components(mask):
H, W = mask.shape; labels = np.zeros_like(mask, dtype=np.int32); cur = 0
for y in range(H):
for x in range(W):
if mask[y, x] and not labels[y, x]:
cur += 1; q = deque([(y, x)]); labels[y, x] = cur
while q:
yy, xx = q.popleft()
for dy, dx in [(-1,0),(1,0),(0,-1),(0,1)]:
ny, nx = yy + dy, xx + dx
if 0 <= ny < H and 0 <= nx < W and mask[ny, nx] and not labels[ny, nx]:
labels[ny, nx] = cur; q.append((ny, nx))
return labels, cur
Tests
m = np.array([[1,1,0,0],[0,1,0,1],[0,0,0,1],[1,0,0,0]], dtype=np.uint8)
lab, n = connected_components(m); assert n == 3
print("CCL OK", lab)
Problem
Find the global threshold that maximises between-class variance on a 256-bin histogram.
import numpy as np
def otsu(img):
h = np.bincount(img.ravel(), minlength=256).astype(np.float64)
p = h / h.sum()
omega = np.cumsum(p); mu = np.cumsum(p * np.arange(256))
mu_t = mu[-1]
sig_b = (mu_t * omega - mu) ** 2 / (omega * (1 - omega) + 1e-12)
return int(np.argmax(sig_b))
Tests
img = np.concatenate([np.full(100, 30), np.full(100, 200)]).astype(np.uint8)
t = otsu(img); assert 30 < t < 200; print("otsu OK", t)
Problem
Encode a binary mask (column-major) as run-length pairs of (start, length).
import numpy as np
def rle_encode(mask):
flat = mask.flatten(order='F').astype(np.uint8)
diffs = np.diff(np.concatenate([[0], flat, [0]]))
starts = np.where(diffs == 1)[0]; ends = np.where(diffs == -1)[0]
return [(int(s), int(e - s)) for s, e in zip(starts, ends)]
def rle_decode(rle, shape):
flat = np.zeros(shape[0] * shape[1], dtype=np.uint8)
for s, l in rle: flat[s:s+l] = 1
return flat.reshape(shape, order='F')
Tests
m = np.array([[1,1,0],[1,0,0]], dtype=np.uint8)
rle = rle_encode(m); back = rle_decode(rle, m.shape)
assert (back == m).all(); print("RLE OK", rle)
Problem
Rasterise a closed polygon onto a 2D mask using the even–odd rule per scanline.
import numpy as np
def polygon_to_mask(poly, H, W):
mask = np.zeros((H, W), dtype=np.uint8)
n = len(poly)
for y in range(H):
xs = []
for i in range(n):
x1, y1 = poly[i]; x2, y2 = poly[(i + 1) % n]
if (y1 <= y < y2) or (y2 <= y < y1):
xs.append(x1 + (y - y1) * (x2 - x1) / (y2 - y1 + 1e-9))
xs.sort()
for k in range(0, len(xs), 2):
x_a, x_b = int(round(xs[k])), int(round(xs[k+1])) if k+1 < len(xs) else 0
mask[y, max(0,x_a):max(0,x_b)] = 1
return mask
Tests
poly = [(2.0,2.0),(8.0,2.0),(8.0,8.0),(2.0,8.0)]
m = polygon_to_mask(poly, 12, 12)
assert m[5, 5] == 1 and m[0, 0] == 0; print("polygon OK")
Problem
Compute per-class IoU and the mean over C classes from prediction and target masks.
def miou(pred, gt, C):
iou = []
for c in range(C):
p = (pred == c); g = (gt == c)
inter = (p & g).sum().float()
uni = (p | g).sum().float().clamp(min=1)
iou.append((inter / uni).item())
return iou, sum(iou) / C
Tests
p = torch.tensor([[0,1,1],[0,1,2]]); g = torch.tensor([[0,1,1],[0,2,2]])
ious, m = miou(p, g, 3)
assert abs(ious[0] - 1.0) < 1e-6; print("mIoU OK", m)
Problem
Given an image I and flow u, v at each pixel, warp I by sampling at (x + u, y + v) via grid_sample.
import torch
import torch.nn.functional as F
def warp_by_flow(img, flow):
# img: (B, C, H, W), flow: (B, 2, H, W) in pixel units
B, _, H, W = img.shape
yy, xx = torch.meshgrid(torch.arange(H, device=img.device).float(),
torch.arange(W, device=img.device).float(), indexing='ij')
grid = torch.stack([xx + flow[:, 0], yy + flow[:, 1]], dim=-1)
grid_norm = torch.zeros_like(grid)
grid_norm[..., 0] = 2.0 * grid[..., 0] / (W - 1) - 1.0
grid_norm[..., 1] = 2.0 * grid[..., 1] / (H - 1) - 1.0
return F.grid_sample(img, grid_norm, align_corners=True)
Tests
img = torch.zeros(1, 1, 4, 4); img[0, 0, 1, 1] = 1.0
flow = torch.zeros(1, 2, 4, 4); flow[0, 0] = 1.0 # shift by +1 in x
out = warp_by_flow(img, flow)
assert out[0, 0, 1, 0].item() > 0.5 # shifted left by 1 in sample location
print("warp flow OK")
Problem
Solve LK in a small window: A⊤Av = −A⊤b where A = [Ix, Iy], b = It.
import torch
import torch.nn.functional as F
def lucas_kanade(I0, I1, k=5):
Gx, Gy, _, _ = sobel(I0)
It = I1 - I0
H, W = I0.shape[-2:]
win = torch.ones(1, 1, k, k, device=I0.device) / (k * k)
Ixx = F.conv2d(Gx * Gx, win, padding=k//2)
Iyy = F.conv2d(Gy * Gy, win, padding=k//2)
Ixy = F.conv2d(Gx * Gy, win, padding=k//2)
Ixt = F.conv2d(Gx * It, win, padding=k//2)
Iyt = F.conv2d(Gy * It, win, padding=k//2)
det = Ixx * Iyy - Ixy ** 2 + 1e-6
u = (-Iyy * Ixt + Ixy * Iyt) / det
v = (Ixy * Ixt - Ixx * Iyt) / det
return torch.cat([u, v], dim=1)
Tests
H, W = 16, 16
I0 = torch.zeros(1, 1, H, W); I0[0, 0, 4:8, 4:8] = 1.0
I1 = torch.zeros(1, 1, H, W); I1[0, 0, 4:8, 5:9] = 1.0 # shift +1 in x
flow = lucas_kanade(I0, I1, k=5)
print("LK OK", flow[0, 0, 6, 6].item(), flow[0, 1, 6, 6].item())
Problem
For each pixel, take arg min over candidate disparities of a precomputed cost volume.
def wta_disparity(cost_vol):
# cost_vol: (D, H, W), lower = better
return cost_vol.argmin(dim=0)
Tests
cv = torch.rand(8, 4, 4); cv[3, 2, 2] = -1.0
d = wta_disparity(cv); assert d[2, 2].item() == 3
print("WTA OK")
Problem
Refine disparity using a parabola fit through neighboring costs: dˆ = d∗+c−1−c+1 2(c−1−2c0+c+1).
def subpixel_refine(cost, d):
D = cost.size(0); H, W = d.shape
out = d.float().clone()
for y in range(H):
for x in range(W):
i = int(d[y, x])
if 1 <= i <= D - 2:
cm, c0, cp = cost[i-1, y, x], cost[i, y, x], cost[i+1, y, x]
denom = 2 * (cm - 2*c0 + cp)
if denom.abs() > 1e-9:
out[y, x] = i + (cm - cp) / denom
return out
Tests
cv = torch.tensor([[1., 0.2, 1.]]).expand(3, 1, 1).clone()
d = wta_disparity(cv)
out = subpixel_refine(cv, d)
print("subpixel OK", out)
Problem
Given a left disparity map dL and right dR, mark pixels where |dL(x) −dR(x −dL(x))| > τ as inconsistent.
import torch
def lr_consistency(dL, dR, tau=1.0):
H, W = dL.shape
valid = torch.zeros_like(dL, dtype=torch.bool)
for y in range(H):
for x in range(W):
xr = int(round(x - dL[y, x].item()))
if 0 <= xr < W:
valid[y, x] = abs(dL[y, x].item() - dR[y, xr].item()) <= tau
return valid
Tests
dL = torch.full((4, 4), 1.0); dR = torch.full((4, 4), 1.0)
v = lr_consistency(dL, dR); assert v.all()
print("LR check OK")
Problem
Implement a 4-state (x, y, x,˙ y˙) constant-velocity Kalman filter for a 2D position observation.
import numpy as np
class KF2D:
def __init__(self, dt=1.0, q=1e-2, r=1.0):
self.F = np.array([[1,0,dt,0],[0,1,0,dt],[0,0,1,0],[0,0,0,1.0]])
self.H = np.array([[1,0,0,0],[0,1,0,0.0]])
self.Q = q * np.eye(4); self.R = r * np.eye(2)
self.x = np.zeros(4); self.P = np.eye(4) * 10
def predict(self):
self.x = self.F @ self.x
self.P = self.F @ self.P @ self.F.T + self.Q
def update(self, z):
y = z - self.H @ self.x
S = self.H @ self.P @ self.H.T + self.R
K = self.P @ self.H.T @ np.linalg.inv(S)
self.x = self.x + K @ y
self.P = (np.eye(4) - K @ self.H) @ self.P
Tests
kf = KF2D(dt=1.0)
for t in range(10):
kf.predict(); kf.update(np.array([t * 1.0 + 0.05*np.random.randn(),
t * 0.5 + 0.05*np.random.randn()]))
assert abs(kf.x[2] - 1.0) < 0.5 # learned ~1 px/frame
print("KF OK", kf.x)
Problem
Match N existing tracks (predicted boxes) to M detections by IoU, using the Hungarian algorithm with a minimum-IoU gate.
from scipy.optimize import linear_sum_assignment
def sort_assoc(tracks, detections, iou_thresh=0.3):
if len(tracks) == 0 or len(detections) == 0: return [], list(range(len(tracks))), list(range(len(detections)))
cost = -iou(torch.as_tensor(tracks).float(), torch.as_tensor(detections).float()).numpy()
r, c = linear_sum_assignment(cost)
matches, ut, ud = [], [], []
for i, j in zip(r, c):
if -cost[i, j] < iou_thresh: ut.append(i); ud.append(j)
else: matches.append((i, j))
ut += [i for i in range(len(tracks)) if i not in r]
ud += [j for j in range(len(detections)) if j not in c]
return matches, ut, ud
def iou(a, b): # local copy of the box-IoU
inter_x1 = torch.maximum(a[:, None, 0], b[None, :, 0])
inter_y1 = torch.maximum(a[:, None, 1], b[None, :, 1])
inter_x2 = torch.minimum(a[:, None, 2], b[None, :, 2])
inter_y2 = torch.minimum(a[:, None, 3], b[None, :, 3])
iw = (inter_x2 - inter_x1).clamp(min=0); ih = (inter_y2 - inter_y1).clamp(min=0)
inter = iw * ih
aA = (a[:, 2] - a[:, 0]) * (a[:, 3] - a[:, 1])
aB = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1])
return inter / (aA[:, None] + aB[None, :] - inter).clamp(min=1e-9)
Tests
T = np.array([[0,0,10,10],[100,100,110,110]], dtype=np.float32)
D = np.array([[1,1,11,11],[200,200,210,210]], dtype=np.float32)
m, ut, ud = sort_assoc(T, D)
assert (0, 0) in m and 1 in ut and 1 in ud; print("SORT OK", m)
Problem
ByteTrack: associate high-score detections first; then re-use unmatched tracks against low-score detections to recover under-occluded targets.
import numpy as np
def bytetrack_step(tracks, dets, scores, high=0.6, low=0.1):
hi = [i for i, s in enumerate(scores) if s >= high]
lo = [i for i, s in enumerate(scores) if low <= s < high]
m1, ut, ud1 = sort_assoc(tracks, dets[hi])
remain = [t for i, t in enumerate(tracks) if i in ut]
if not remain or not lo: return m1
m2, _, _ = sort_assoc(np.asarray(remain), dets[lo], iou_thresh=0.5)
return m1 + m2
Tests
T = np.array([[0,0,10,10]], dtype=np.float32)
D = np.array([[1,1,11,11],[2,2,12,12]], dtype=np.float32)
S = [0.9, 0.2]
m = bytetrack_step(T, D, S)
assert (0, 0) in m; print("byte OK", m)
Problem
Implement Hamilton-product quaternion multiplication (q = q1⊗q2) using the standard (w, x, y, z) convention.
import torch
def quat_mul(q1, q2):
w1, x1, y1, z1 = q1
w2, x2, y2, z2 = q2
return torch.stack([
w1*w2 - x1*x2 - y1*y2 - z1*z2,
w1*x2 + x1*w2 + y1*z2 - z1*y2,
w1*y2 - x1*z2 + y1*w2 + z1*x2,
w1*z2 + x1*y2 - y1*x2 + z1*w2])
Tests
q1 = torch.tensor([1., 0, 0, 0]) # identity
q2 = torch.tensor([math.cos(0.5), math.sin(0.5), 0, 0]) # rot 1 rad about x
q = quat_mul(q1, q2)
assert torch.allclose(q, q2, atol=1e-6); print("quat OK")
Advanced implementation. Batched over arbitrary leading dimensions via unbind/stack — one call for a whole trajectory of rotations. Verified against scipy.spatial.transform.Rotation composition (accounting for the q/-q double cover).
import torch
def quat_mul(q1, q2):
# (..., 4) wxyz convention, fully broadcast
w1, x1, y1, z1 = q1.unbind(-1); w2, x2, y2, z2 = q2.unbind(-1)
return torch.stack([
w1*w2 - x1*x2 - y1*y2 - z1*z2,
w1*x2 + x1*w2 + y1*z2 - z1*y2,
w1*y2 - x1*z2 + y1*w2 + z1*x2,
w1*z2 + x1*y2 - y1*x2 + z1*w2], -1)
Problem
Convert a 3×3 rotation matrix to a unit quaternion (w, x, y, z).
import numpy as np
import math
def mat_to_quat(R):
tr = R[0, 0] + R[1, 1] + R[2, 2]
if tr > 0:
S = math.sqrt(tr + 1.0) * 2
w = 0.25 * S
x = (R[2, 1] - R[1, 2]) / S
y = (R[0, 2] - R[2, 0]) / S
z = (R[1, 0] - R[0, 1]) / S
elif (R[0,0] > R[1,1]) and (R[0,0] > R[2,2]):
S = math.sqrt(1.0 + R[0,0] - R[1,1] - R[2,2]) * 2
w = (R[2,1] - R[1,2]) / S; x = 0.25 * S
y = (R[0,1] + R[1,0]) / S; z = (R[0,2] + R[2,0]) / S
elif R[1,1] > R[2,2]:
S = math.sqrt(1.0 + R[1,1] - R[0,0] - R[2,2]) * 2
w = (R[0,2] - R[2,0]) / S; x = (R[0,1] + R[1,0]) / S
y = 0.25 * S; z = (R[1,2] + R[2,1]) / S
else:
S = math.sqrt(1.0 + R[2,2] - R[0,0] - R[1,1]) * 2
w = (R[1,0] - R[0,1]) / S; x = (R[0,2] + R[2,0]) / S
y = (R[1,2] + R[2,1]) / S; z = 0.25 * S
return np.array([w, x, y, z])
Tests
R = np.eye(3)
q = mat_to_quat(R); assert np.allclose(q, [1,0,0,0])
print("mat_to_quat OK")
Problem
Map a twist ξ = (ρ, ϕ) ∈ R6 (translational + rotational) to a homogeneous transform.
import numpy as np
def hat(omega):
return np.array([[0, -omega[2], omega[1]],
[omega[2], 0, -omega[0]],
[-omega[1], omega[0], 0.0]])
def se3_exp(xi):
rho, phi = xi[:3], xi[3:]
theta = np.linalg.norm(phi)
Phi = hat(phi)
if theta < 1e-9:
R = np.eye(3) + Phi
J = np.eye(3) + 0.5 * Phi
else:
R = np.eye(3) + np.sin(theta) / theta * Phi + (1 - np.cos(theta)) / theta**2 * Phi @ Phi
J = np.eye(3) + (1 - np.cos(theta)) / theta**2 * Phi + (theta - np.sin(theta)) / theta**3 * Phi @ Phi
T = np.eye(4); T[:3, :3] = R; T[:3, 3] = J @ rho
return T
Tests
T = se3_exp(np.array([1, 2, 3, 0, 0, 0.0]))
assert np.allclose(T[:3, 3], [1, 2, 3]); print("se3 exp OK")
Problem
Apply a rotation R and translation t to a point cloud P ∈ RN×3.
def transform_pc(P, R, t):
return P @ R.T + t
Tests
P = np.eye(3)
T = transform_pc(P, np.eye(3), np.array([1, 2, 3.]))
assert np.allclose(T - P, np.array([1, 2, 3.])); print("transform OK")
Problem
One ICP iteration: find nearest-neighbor pairs, solve the closed-form rigid transform via SVD (Kabsch).
import numpy as np
def icp_step(P, Q):
# nearest neighbor: brute force
d = ((P[:, None] - Q[None]) ** 2).sum(-1)
idx = d.argmin(axis=1)
Qm = Q[idx]
cP = P.mean(0); cQ = Qm.mean(0)
H = (P - cP).T @ (Qm - cQ)
U, _, Vt = np.linalg.svd(H)
R = Vt.T @ U.T
if np.linalg.det(R) < 0: Vt[-1] *= -1; R = Vt.T @ U.T
t = cQ - R @ cP
return R, t
Tests
np.random.seed(0); P = np.random.randn(50, 3)
R_true = np.array([[0,-1,0],[1,0,0],[0,0,1.0]])
t_true = np.array([0.5, -0.3, 1.0])
Q = P @ R_true.T + t_true
R_e, t_e = icp_step(P, Q)
assert np.allclose(R_e, R_true, atol=1e-6); assert np.allclose(t_e, t_true, atol=1e-6)
print("ICP OK")
Problem
Symmetric Chamfer distance between two point clouds.
def chamfer(P, Q):
d = ((P[:, None] - Q[None]) ** 2).sum(-1)
return d.min(axis=1).mean() + d.min(axis=0).mean()
Tests
P = np.random.randn(50, 3); Q = P + 0.01 * np.random.randn(50, 3)
assert chamfer(P, Q) < 0.01; print("chamfer OK")
Advanced implementation. The full (N, M) distance matrix is the memory bottleneck at point-cloud scale (100k x 100k = 40 GB). Tile one side; results are exact. Same chunk-the-batch pattern as FlashAttention tiling. Verified equal to the full-matrix version (1e-5).
import torch
def chamfer_chunked(A, B, chunk=4096):
mins_a = torch.cat([torch.cdist(A[i:i+chunk], B).min(1).values
for i in range(0, len(A), chunk)])
mins_b = torch.cat([torch.cdist(B[i:i+chunk], A).min(1).values
for i in range(0, len(B), chunk)])
return (mins_a ** 2).mean() + (mins_b ** 2).mean()
Problem
Voxelize a point cloud into a V ×V ×V occupancy grid.
import numpy as np
def voxelize(P, V=32, bounds=(-1, 1)):
lo, hi = bounds
idx = np.floor((P - lo) / (hi - lo) * V).astype(int).clip(0, V - 1)
grid = np.zeros((V, V, V), dtype=np.uint8)
grid[idx[:, 0], idx[:, 1], idx[:, 2]] = 1
return grid
Tests
P = np.random.uniform(-1, 1, (1000, 3))
g = voxelize(P, V=8); assert g.sum() > 0 and g.shape == (8, 8, 8)
print("voxelize OK")
Problem
Pick K points from a cloud by greedy farthest-point sampling.
import numpy as np
def fps(P, K, rng=None):
rng = rng or np.random.default_rng(0)
N = len(P); idx = [int(rng.integers(N))]
dist = np.full(N, np.inf)
for _ in range(K - 1):
d = ((P - P[idx[-1]]) ** 2).sum(-1)
dist = np.minimum(dist, d)
idx.append(int(np.argmax(dist)))
return idx
Tests
P = np.random.randn(200, 3)
idx = fps(P, K=10); assert len(set(idx)) == 10; print("FPS OK")
Advanced implementation. Keep a running min-distance-to-selected-set: each new pick needs distances to the NEWEST point only — O(nk) instead of the naive O(n·k²)-ish recompute against the whole selected set. This is how PointNet++ CUDA kernels do it. Verified: identical pick sequence to the naive version.
import torch
def fps_fast(pts, k, start=0):
sel = [start]
d2 = ((pts - pts[start]) ** 2).sum(1) # running min-dist to selected set
for _ in range(k - 1):
nxt = int(d2.argmax())
sel.append(nxt)
d2 = torch.minimum(d2, ((pts - pts[nxt]) ** 2).sum(1)) # only the NEW point
return sel
Problem
Split a (B, C, H, W) image into non-overlapping p×p patches and reshape to (B, N, p2C).
def patchify(img, p=16):
B, C, H, W = img.shape
assert H % p == 0 and W % p == 0
x = img.unfold(2, p, p).unfold(3, p, p) # (B, C, H/p, W/p, p, p)
x = x.permute(0, 2, 3, 1, 4, 5).contiguous() # (B, H/p, W/p, C, p, p)
return x.view(B, -1, C * p * p)
Tests
x = torch.randn(2, 3, 32, 32); y = patchify(x, p=16)
assert y.shape == (2, 4, 3 * 16 * 16); print("patchify OK")
Problem
Prepend a learnable CLS token to a sequence of patch tokens and add a learnable PE.
import torch
import torch.nn as nn
class CLSPlusPos(nn.Module):
def __init__(self, dim, n_tokens):
super().__init__()
self.cls = nn.Parameter(torch.zeros(1, 1, dim))
self.pe = nn.Parameter(torch.zeros(1, n_tokens + 1, dim))
nn.init.trunc_normal_(self.cls, std=0.02)
nn.init.trunc_normal_(self.pe, std=0.02)
def forward(self, x):
B = x.size(0)
cls = self.cls.expand(B, -1, -1)
return torch.cat([cls, x], dim=1) + self.pe
Tests
m = CLSPlusPos(64, n_tokens=10); y = m(torch.randn(2, 10, 64))
assert y.shape == (2, 11, 64); print("CLS+PE OK")
Problem
Split a (B, H, W, C) tensor into windows of size w and reconstruct the original layout.
def window_partition(x, w):
B, H, W, C = x.shape
x = x.view(B, H // w, w, W // w, w, C)
return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, w, w, C)
def window_reverse(windows, w, H, W):
B = int(windows.size(0) / (H * W / (w * w)))
x = windows.view(B, H // w, W // w, w, w, -1)
return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
Tests
x = torch.randn(2, 8, 8, 16)
w = window_partition(x, 4)
back = window_reverse(w, 4, 8, 8)
assert torch.allclose(back, x); print("swin window OK")
Problem
Build a Swin-style relative positional bias indexed by (∆y, ∆x) within a window.
import torch
import torch.nn as nn
class RelPosBias(nn.Module):
def __init__(self, w, h):
super().__init__()
self.table = nn.Parameter(torch.zeros((2 * h - 1) * (2 * w - 1)))
coords = torch.stack(torch.meshgrid(torch.arange(h), torch.arange(w), indexing='ij')) # (2,h,w)
flat = coords.flatten(1) # (2, hw)
rel = flat[:, :, None] - flat[:, None, :] # (2, hw, hw)
rel = rel.permute(1, 2, 0).contiguous()
rel[:, :, 0] += h - 1; rel[:, :, 1] += w - 1
rel[:, :, 0] *= 2 * w - 1
idx = rel.sum(-1)
self.register_buffer('idx', idx)
def forward(self):
return self.table[self.idx]
Tests
m = RelPosBias(4, 4); b = m()
assert b.shape == (16, 16); print("relpos OK")
Problem
Peak signal-to-noise ratio for an 8-bit image (in dB).
import torch
def psnr(x, y, max_val=1.0):
mse = ((x - y) ** 2).mean()
return 10 * torch.log10(max_val ** 2 / mse.clamp(min=1e-12))
Tests
x = torch.rand(3, 8, 8); y = x + 1e-3 * torch.randn_like(x)
assert psnr(x, y) > 40; print("psnr OK")
Problem
Implement a basic single-scale SSIM with 11×11 Gaussian window.
import torch.nn.functional as F
def ssim(x, y, win=11, C1=0.01**2, C2=0.03**2):
g = gaussian_kernel_1d(1.5, k=win).to(x.device)
Cn = x.size(1)
g = g.view(1, 1, 1, win).expand(Cn, 1, 1, win)
pad = win // 2
def filt(z): return F.conv2d(F.conv2d(z, g, padding=(0, pad), groups=Cn),
g.transpose(2, 3), padding=(pad, 0), groups=Cn)
mu_x = filt(x); mu_y = filt(y)
mu_x2, mu_y2, mu_xy = mu_x ** 2, mu_y ** 2, mu_x * mu_y
sx = filt(x * x) - mu_x2; sy = filt(y * y) - mu_y2; sxy = filt(x * y) - mu_xy
return (((2 * mu_xy + C1) * (2 * sxy + C2)) /
((mu_x2 + mu_y2 + C1) * (sx + sy + C2))).mean()
Tests
x = torch.rand(1, 1, 32, 32)
assert abs(ssim(x, x).item() - 1.0) < 1e-4; print("ssim OK")
Problem
Compute average precision for a single class given matched / unmatched predictions sorted by score.
import numpy as np
def average_precision(matches, n_gt):
# matches: list of 1 (TP) or 0 (FP) sorted by score descending
tp = np.cumsum(matches); fp = np.cumsum(1 - np.asarray(matches))
rec = tp / max(1, n_gt); prec = tp / np.maximum(tp + fp, 1e-9)
# 11-point interpolated AP
ap = 0.0
for r in np.linspace(0, 1, 11):
p = prec[rec >= r].max() if (rec >= r).any() else 0.0
ap += p / 11
return float(ap)
Tests
ap = average_precision([1,1,0,1], n_gt=4)
assert 0.5 < ap <= 1.0; print("AP OK", ap)
Problem
Given prediction scores, target labels (binary), and a sweep of thresholds, return precision and recall arrays.
import numpy as np
def pr_curve(scores, targets):
order = np.argsort(-scores)
s = scores[order]; t = targets[order]
tp = np.cumsum(t); fp = np.cumsum(1 - t)
prec = tp / np.maximum(tp + fp, 1)
rec = tp / max(1, t.sum())
return prec, rec, s
Tests
s = np.array([0.9, 0.4, 0.7, 0.2]); t = np.array([1, 0, 1, 0])
p, r, _ = pr_curve(s, t); assert r[-1] == 1.0; print("PR OK")
Problem
Probability of Correct Keypoint at threshold α: a prediction is correct if its distance to the GT is < α · d where d is the bounding-box diagonal.
def pck(pred, gt, bbox_diag, alpha=0.05):
d = (pred - gt).norm(dim=-1)
return (d < alpha * bbox_diag).float().mean().item()
Tests
gt = torch.zeros(10, 2); pr = gt + 0.001
assert pck(pr, gt, bbox_diag=1.0) == 1.0; print("PCK OK")
Problem
Compute box AP averaged over IoU thresholds 0.50, 0.55, . . . , 0.95 for one class.
import numpy as np
import torch
def coco_ap(preds, scores, gts):
# preds: (N,4) sorted by scores desc; gts: (M, 4)
aps = []
for thr in np.linspace(0.5, 0.95, 10):
matched_gt = set()
m = []
for i, p in enumerate(preds):
if len(gts) == 0:
m.append(0); continue
ious = iou(torch.as_tensor(p)[None].float(), torch.as_tensor(gts).float()).numpy()[0]
j = ious.argmax()
if ious[j] >= thr and j not in matched_gt:
matched_gt.add(j); m.append(1)
else:
m.append(0)
aps.append(average_precision(m, len(gts)))
return float(np.mean(aps))
Tests
preds = np.array([[0,0,10,10],[20,20,30,30]], dtype=np.float32)
scores = np.array([0.9, 0.8])
gts = np.array([[0,0,10,10]], dtype=np.float32)
assert coco_ap(preds, scores, gts) > 0; print("coco AP OK")
Problem
Implement random_crop(img, h, w) and random_hflip(img, p=0.5).
import numpy as np
import torch
def random_crop(img, h, w, rng=None):
rng = rng or np.random
H, W = img.shape[-2:]
y = rng.randint(0, H - h + 1); x = rng.randint(0, W - w + 1)
return img[..., y:y+h, x:x+w], (y, x)
def random_hflip(img, p=0.5, rng=None):
rng = rng or np.random
if rng.random() < p: return torch.flip(img, dims=[-1]), True
return img, False
Tests
x = torch.arange(16).view(1, 1, 4, 4).float()
o, _ = random_crop(x, 2, 2); assert o.shape[-2:] == (2, 2)
print("crop/flip OK")
Problem
Apply random brightness, contrast, saturation, hue jitter to an RGB image.
import numpy as np
def color_jitter(img, b=0.2, c=0.2, s=0.2, h=0.05, rng=None):
rng = rng or np.random
img = img * (1 + rng.uniform(-b, b))
mean = img.mean(dim=(-1, -2), keepdim=True)
img = (img - mean) * (1 + rng.uniform(-c, c)) + mean
gray = rgb_to_gray(img)
img = (img - gray) * (1 + rng.uniform(-s, s)) + gray
hsv = rgb_to_hsv(img.clamp(0, 1))
hsv[..., 0, :, :] = (hsv[..., 0, :, :] + rng.uniform(-h, h)) % 1.0
return hsv_to_rgb(hsv).clamp(0, 1)
Tests
out = color_jitter(torch.rand(1, 3, 8, 8))
assert out.shape == (1, 3, 8, 8); print("jitter OK")
Problem
Erase a random rectangular region of an image with probability p, filled with random noise.
import numpy as np
import math
import torch
def random_erase(img, p=0.5, sl=0.02, sh=0.2, rng=None):
rng = rng or np.random
if rng.random() > p: return img
C, H, W = img.shape[-3:]
s = rng.uniform(sl, sh) * H * W
r = rng.uniform(0.3, 1/0.3)
h = int(round(math.sqrt(s * r))); w = int(round(math.sqrt(s / r)))
if h >= H or w >= W: return img
y = rng.randint(0, H - h); x = rng.randint(0, W - w)
img = img.clone()
img[..., y:y+h, x:x+w] = torch.randn_like(img[..., y:y+h, x:x+w])
return img
Tests
out = random_erase(torch.zeros(3, 8, 8), p=1.0); assert out.abs().sum() > 0
print("erase OK")
Problem
Concatenate four images of size (C, H, W) into a 2H × 2W mosaic; return the corresponding box transforms.
import torch
def mosaic(imgs):
# imgs: list of 4 tensors (C, H, W) of equal size
C, H, W = imgs[0].shape
out = torch.zeros(C, 2 * H, 2 * W)
out[:, :H, :W] = imgs[0]; out[:, :H, W:] = imgs[1]
out[:, H:, :W] = imgs[2]; out[:, H:, W:] = imgs[3]
return out
Tests
m = mosaic([torch.full((3, 4, 4), float(i)) for i in range(4)])
assert m[:, 0, 0].mean() == 0 and m[:, 0, 5].mean() == 1
print("mosaic OK")
Problem
Sample N random ops from a list, each parameterised by magnitude M, and apply sequentially.
import numpy as np
def randaugment(img, ops, N=2, M=10, rng=None):
rng = rng or np.random
chosen = rng.choice(len(ops), size=N, replace=False)
for k in chosen:
img = ops[k](img, M)
return img
# helpers (simplified)
def op_brightness(img, M): return (img * (1 + 0.05 * M)).clamp(0, 1)
def op_contrast(img, M):
mean = img.mean(); return ((img - mean) * (1 + 0.05 * M) + mean).clamp(0, 1)
Tests
x = torch.rand(3, 8, 8); ops = [op_brightness, op_contrast]
y = randaugment(x, ops, N=2, M=5); assert y.shape == x.shape
print("randaugment OK")
Problem
Implement Grad-CAM for a CNN: compute gradients of a target logit w.r.t. the activations of a chosen conv layer; weight channels by global-average-pooled gradients; ReLU, normalize.
def grad_cam(model, x, target_layer, class_idx):
feats = []; grads = []
h1 = target_layer.register_forward_hook(lambda m, i, o: feats.append(o))
h2 = target_layer.register_full_backward_hook(lambda m, gi, go: grads.append(go[0]))
out = model(x); model.zero_grad()
out[0, class_idx].backward()
A = feats[0]; dA = grads[0]
w = dA.mean(dim=(2, 3), keepdim=True)
cam = (w * A).sum(dim=1, keepdim=True).relu()
cam = cam / (cam.amax(dim=(2, 3), keepdim=True) + 1e-8)
h1.remove(); h2.remove()
return cam
Tests
class TinyNet(nn.Module):
def __init__(self):
super().__init__()
self.c = nn.Conv2d(3, 4, 3, padding=1)
self.f = nn.Linear(4, 2)
def forward(self, x):
h = self.c(x); self.last = h
return self.f(h.mean(dim=(2, 3)))
m = TinyNet(); cam = grad_cam(m, torch.rand(1, 3, 16, 16, requires_grad=True), m.c, 0)
assert cam.shape == (1, 1, 16, 16); print("grad-cam OK")
Problem
∫10 ∂xf(x′ + α(x −x′))dα. Approximate with S steps. IG(x) = (x −x′)
import torch
def integrated_gradients(model, x, baseline=None, target=0, S=50):
if baseline is None: baseline = torch.zeros_like(x)
accum = torch.zeros_like(x)
for s in range(1, S + 1):
a = s / S
z = (baseline + a * (x - baseline)).requires_grad_(True)
out = model(z)
out[:, target].sum().backward()
accum += z.grad
return (x - baseline) * accum / S
Tests
class IDM(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 2, bias=False)
def forward(self, x): return self.l(x)
m = IDM(); m.l.weight.data = torch.eye(4)[:2]
ig = integrated_gradients(m, torch.tensor([[1., 1., 0., 0.]]), target=0, S=10)
assert abs(ig[0, 0] - 1.0) < 1e-3; print("IG OK")
Problem
Smooth a saliency map by averaging gradients over S noisy copies of the input.
import torch
def smoothgrad(model, x, target=0, sigma=0.1, S=20):
accum = torch.zeros_like(x)
for _ in range(S):
z = (x + sigma * torch.randn_like(x)).requires_grad_(True)
out = model(z)
out[:, target].sum().backward()
accum += z.grad
return accum / S
Tests
class IDM(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 2, bias=False)
def forward(self, x): return self.l(x)
m = IDM(); m.l.weight.data = torch.eye(4)[:2]
out = smoothgrad(m, torch.rand(1, 4), target=0, sigma=0.01, S=5)
assert out.shape == (1, 4); print("smoothgrad OK")
Problem
Given a list of (kernel, stride, padding) for a sequential conv stack, compute the input receptive field of the last layer.
def receptive_field(layers):
rf = 1; jump = 1
for k, s, p in layers:
rf = rf + (k - 1) * jump
jump *= s
return rf
Tests
assert receptive_field([(3, 1, 1), (3, 1, 1), (3, 1, 1)]) == 7
assert receptive_field([(7, 2, 3), (3, 2, 1), (3, 1, 1)]) == 19
print("RF OK")
Problem
For each pixel (u, v), generate a ray in world space given camera intrinsics K and pose Twc.
import numpy as np
def gen_rays(H, W, K, T_wc):
# T_wc: 4x4 camera-to-world
yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
inv_K = np.linalg.inv(K)
pix = np.stack([xx, yy, np.ones_like(xx)], axis=-1).reshape(-1, 3)
cam_dirs = pix @ inv_K.T
cam_dirs = cam_dirs / np.linalg.norm(cam_dirs, axis=-1, keepdims=True)
R = T_wc[:3, :3]; t = T_wc[:3, 3]
world_dirs = cam_dirs @ R.T
origins = np.broadcast_to(t, world_dirs.shape).copy()
return origins.reshape(H, W, 3), world_dirs.reshape(H, W, 3)
Tests
K = np.array([[100,0,32],[0,100,32],[0,0,1.0]])
o, d = gen_rays(64, 64, K, np.eye(4))
assert d[32, 32, 2] > 0.99; print("gen_rays OK")
Problem
Compute the NeRF positional encoding γ(p) = (sin(2lπp), cos(2lπp))l=0,...,L−1.
import math
import torch
def pos_enc(p, L=10):
freqs = 2.0 ** torch.arange(L, device=p.device) * math.pi
out = [p]
for f in freqs: out += [torch.sin(f * p), torch.cos(f * p)]
return torch.cat(out, dim=-1)
Tests
p = torch.rand(4, 3)
out = pos_enc(p, L=4); assert out.shape == (4, 3 + 3 * 2 * 4)
print("posenc OK")
Problem
i Ti(1 −e−σiδi)ci.Integrate σ, c along sorted samples to produce pixel color and accumulated alpha: C = Σ
import torch
def volume_render(sigma, color, deltas):
# sigma: (N,), color: (N, 3), deltas: (N,)
alpha = 1 - torch.exp(-sigma * deltas)
T = torch.cumprod(torch.cat([torch.ones(1), 1 - alpha + 1e-10]), dim=0)[:-1]
w = T * alpha
return (w[:, None] * color).sum(dim=0), w.sum()
Tests
sigma = torch.tensor([10.0, 0.0, 0.0]); color = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1.]])
delt = torch.tensor([0.1, 0.1, 0.1])
c, a = volume_render(sigma, color, delt)
assert (c[0] > 0.5).item() and a > 0.5; print("volrender OK")
Problem
Given coarse weights along a ray, draw N fine samples by inverse-CDF sampling (NeRF sample_pdf).
import torch
def sample_pdf(bins, weights, N, det=True):
# bins: (B+1,), weights: (B,)
weights = weights + 1e-5
pdf = weights / weights.sum()
cdf = torch.cumsum(pdf, dim=0)
cdf = torch.cat([torch.zeros(1, device=cdf.device), cdf])
u = torch.linspace(0, 1, N, device=cdf.device) if det else torch.rand(N, device=cdf.device)
inds = torch.searchsorted(cdf, u.contiguous(), right=True) - 1
inds = inds.clamp(0, len(bins) - 2)
cdf_g = cdf[inds]; bins_g = bins[inds]
denom = (cdf[inds + 1] - cdf_g).clamp(min=1e-9)
t = (u - cdf_g) / denom
return bins_g + t * (bins[inds + 1] - bins[inds])
Tests
bins = torch.linspace(0, 1, 11)
w = torch.zeros(10); w[5] = 1.0
s = sample_pdf(bins, w, 8)
assert (s > bins[5] - 0.01).all() and (s < bins[6] + 0.01).all()
print("sample_pdf OK")
Problem
Evaluate the 9 SH basis functions (ℓ= 0..2) at a unit direction (x, y, z).
import torch
def sh2_basis(d):
x, y, z = d[..., 0], d[..., 1], d[..., 2]
out = torch.stack([
torch.full_like(x, 0.282095), # Y00
-0.488603 * y, 0.488603 * z, -0.488603 * x, # Y1m
1.092548 * x * y, -1.092548 * y * z,
0.315392 * (3 * z * z - 1),
-1.092548 * x * z, 0.546274 * (x * x - y * y) # Y2m
], dim=-1)
return out
Tests
d = torch.tensor([[0., 0., 1.0]])
out = sh2_basis(d)
assert out.shape == (1, 9); print("sh2 OK")
Problem
Project a 3D Gaussian center via K[R|t] and approximate its 2D covariance: Σ2D ≈ J W Σ W ⊤J⊤where J is the Jacobian of the perspective projection at µc.
import numpy as np
def project_gaussian(mu, Sigma, K, R, t):
mu_c = R @ mu + t # camera frame (3,)
f_x, f_y = K[0, 0], K[1, 1]
z = mu_c[2]
J = np.array([[f_x / z, 0, -f_x * mu_c[0] / z**2],
[0, f_y / z, -f_y * mu_c[1] / z**2]])
W = R
Sigma2D = J @ W @ Sigma @ W.T @ J.T
uv = K[:2, :3] @ (mu_c / z)
return uv, Sigma2D
Tests
K = np.array([[100,0,32],[0,100,32],[0,0,1.0]])
uv, S = project_gaussian(np.array([0,0,1.0]), 0.01*np.eye(3), K, np.eye(3), np.zeros(3))
assert np.allclose(uv, [32, 32]) and S.shape == (2, 2)
print("project_gauss OK")
Problem
Given K, R, t and a set of 3D–2D correspondences, compute mean reprojection error in pixels.
import numpy as np
def reproj_err(K, R, t, X, x):
Xc = X @ R.T + t
p = Xc @ K.T
pred = p[:, :2] / p[:, 2:3]
return np.linalg.norm(pred - x, axis=-1).mean()
Tests
K = np.array([[100,0,32],[0,100,32],[0,0,1.0]])
X = np.array([[0,0,1.0],[1,0,1],[0,1,1.0]])
x = (X @ K.T)[:, :2] / (X @ K.T)[:, 2:3]
e = reproj_err(K, np.eye(3), np.zeros(3), X, x)
assert e < 1e-9; print("reproj OK")
Problem
Build an L-level Gaussian pyramid by iterative blur + 2× downsample.
import torch.nn.functional as F
def gaussian_pyramid(img, L=4, sigma=1.0):
pyr = [img]
for _ in range(L - 1):
blurred = gaussian_blur(pyr[-1], sigma)
pyr.append(F.avg_pool2d(blurred, 2))
return pyr
Tests
x = torch.rand(1, 3, 32, 32); pyr = gaussian_pyramid(x, L=3)
assert pyr[-1].shape[-2:] == (8, 8); print("pyramid OK")
Problem
Implement a single-bin ROI Align: bilinearly sample the feature map at four sub-bin points and average.
def roi_align(feat, boxes, out_size=7):
# feat: (1, C, H, W); boxes: (N, 4) in feature-map coords
return torchvision.ops.roi_align(feat, [boxes.float()], output_size=out_size, aligned=True)
Tests
import torchvision
x = torch.randn(1, 3, 16, 16)
b = torch.tensor([[2., 2., 8., 8.]])
out = roi_align(x, b, out_size=4)
assert out.shape == (1, 3, 4, 4); print("ROIAlign OK")
Problem
Convert a heatmap to a continuous keypoint location via spatial soft-argmax.
import torch
def soft_argmax(heat, temperature=1.0):
# heat: (B, K, H, W)
B, K, H, W = heat.shape
flat = (heat / temperature).flatten(2).softmax(dim=-1).view(B, K, H, W)
yy, xx = torch.meshgrid(torch.arange(H, device=heat.device).float(),
torch.arange(W, device=heat.device).float(), indexing='ij')
x = (flat * xx).sum(dim=(2, 3))
y = (flat * yy).sum(dim=(2, 3))
return torch.stack([x, y], dim=-1)
Tests
h = torch.zeros(1, 1, 8, 8); h[0, 0, 5, 3] = 10.0
kp = soft_argmax(h, temperature=0.1)
assert torch.allclose(kp, torch.tensor([[[3., 5.]]]), atol=0.1)
print("softargmax OK")
Problem
Given many noisy 2D–3D correspondences with confidence scores wi, minimize the weighted reprojection error via Gauss–Newton (one step).
import numpy as np
def gn_step_pnp(K, R, t, X, x, w):
Xc = X @ R.T + t
p = Xc @ K.T
pred = p[:, :2] / p[:, 2:3]
res = (pred - x).reshape(-1) # 2N
# Jacobian wrt translation only (small step)
J = np.zeros((2 * len(X), 3))
fx, fy = K[0, 0], K[1, 1]
for i, Xi in enumerate(Xc):
z = Xi[2]
J[2*i ] = [fx / z, 0, -fx * Xi[0] / z**2]
J[2*i+1] = [0, fy / z, -fy * Xi[1] / z**2]
W = np.diag(np.repeat(w, 2))
dt = -np.linalg.solve(J.T @ W @ J + 1e-6 * np.eye(3), J.T @ W @ res)
return t + dt
Tests
K = np.array([[100,0,32],[0,100,32],[0,0,1.0]])
X = np.array([[0,0,1.],[1,0,1.],[0,1,1.]])
x = np.array([[32,32],[132,32],[32,132.]])
t_new = gn_step_pnp(K, np.eye(3), np.array([0.1,0.1,0.0]), X, x, np.ones(3))
assert np.linalg.norm(t_new) < 0.5; print("GN PnP OK")
Notes
Survival tactics for live CV coding:
Always state image conventions: (B, C, H, W) for PyTorch, (H, W, C) for OpenCV / NumPy. Get the order right before any indexing.
Coordinate convention: pixel (u, v) corresponds to (xcolumn, yrow). When sampling from a feature map, the top-left pixel is (0, 0), and align_corners=True vs. False matters for pixel-perfect tasks.
Geometry: write down the K[R|t] projection. Verify with a single point projection sanity check (e.g. a point at the principal axis should project to the principal point).
For RANSAC: always state the minimal sample size (line: 2, homography: 4, fundamental: 8, PnP: 6) and the inlier threshold in pixel units.
For boxes: re-derive the IoU formula by drawing two boxes; include the (x2 - x1).clamp(min=0) step. This is the most common bug source.
For mAP: distinguish 11-point Pascal AP, COCO AP@[0.5:0.95], and per-class vs. averaged.
For ViTs and Swin: keep the (B, H, W, C) vs. (B, N, C) permutations explicit; the window partition is a permute + view, easy to get wrong.
For NeRF / 3DGS: derive volume rendering from the discrete alpha compositing, not the other way round. Mention that the EWA Jacobian is the load-bearing approximation.
Always end every solution with at least one assertion that catches the obvious bug (shape, sum-to-1, sign of GIoU, identity invariance, etc.).