Week 13 — Deep Learning for Computer Vision in Robotics Day 6 of 7
Today’s Goal
In this session we build a small, fully self-contained PyTorch pipeline that estimates the 6D pose (3D rotation + 3D translation) of a known object from a single RGB image. We train a convolutional neural network from scratch on a synthetic dataset we generate ourselves, so you do not need LINEMOD, YCB-Video, or any external download. By the end you will be able to:
- Generate paired (image, pose) training data with a simple pinhole camera model.
- Build a CNN that regresses a 3D translation vector and a 6D rotation representation (Zhou et al., 2019) that converts cleanly to a rotation matrix.
- Train end-to-end with a combined translation + rotation loss and watch both terms decrease.
- Decode predictions back to a pose, project the object’s 3D corners into the image, and measure the reprojection error in pixels.
The model is deliberately small (~1.2M parameters) so it trains in a couple of minutes on CPU.
Prerequisites
- Python 3.10 or newer.
- Install the exact dependencies:
pip install torch==2.3.1 numpy==1.26.4 matplotlib==3.9.0
No torchvision, no OpenCV, no SciPy. We use NumPy for geometry (rotation matrices, projection) and Matplotlib only for the final visualization step. All code runs on CPU; if you have CUDA the script will use it automatically.
Step-by-Step Implementation
Step 1: Setup and Imports
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
SEED = 42
np.random.seed(SEED)
torch.manual_seed(SEED)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMG_SIZE = 64 # images are 64x64 grayscale
FX = FY = 120.0 # focal length in pixels
CX = CY = IMG_SIZE / 2 # principal point
# Object: a unit cube centered at the origin, half-extent 0.5 m
CUBE_CORNERS = np.array([
[-0.5, -0.5, -0.5], [ 0.5, -0.5, -0.5], [ 0.5, 0.5, -0.5], [-0.5, 0.5, -0.5],
[-0.5, -0.5, 0.5], [ 0.5, -0.5, 0.5], [ 0.5, 0.5, 0.5], [-0.5, 0.5, 0.5],
], dtype=np.float64)
def euler_to_matrix(rx, ry, rz):
"""ZYX intrinsic Euler angles -> 3x3 rotation matrix."""
cx, sx = math.cos(rx), math.sin(rx)
cy, sy = math.cos(ry), math.sin(ry)
cz, sz = math.cos(rz), math.sin(rz)
Rx = np.array([[1,0,0],[0,cx,-sx],[0,sx,cx]])
Ry = np.array([[cy,0,sy],[0,1,0],[-sy,0,cy]])
Rz = np.array([[cz,-sz,0],[sz,cz,0],[0,0,1]])
return Rz @ Ry @ Rx
def matrix_to_6d(R):
"""Rotation matrix -> 6D representation (first two columns)."""
return R[:, :2].reshape(-1).astype(np.float32)
def sixd_to_matrix(v):
"""6D representation -> rotation matrix via Gram-Schmidt (Zhou et al. 2019)."""
a1, a2 = v[:3], v[3:]
b1 = a1 / (np.linalg.norm(a1) + 1e-8)
b2 = a2 - np.dot(b1, a2) * b1
b2 = b2 / (np.linalg.norm(b2) + 1e-8)
b3 = np.cross(b1, b2)
return np.stack([b1, b2, b3], axis=1)
def project(points_cam):
"""Pinhole projection: Nx3 camera-frame points -> Nx2 pixel coordinates."""
z = np.clip(points_cam[:, 2], 1e-3, None)
u = FX * points_cam[:, 0] / z + CX
v = FY * points_cam[:, 1] / z + CY
return np.stack([u, v], axis=1)
The key design decision is the 6D rotation representation. Regressing a 3x3 matrix directly with an L2 loss gives a matrix that is not orthonormal, and regressing a quaternion forces you to pick a sign convention that flips discontinuously. The 6D form (two learned columns, orthonormalized by Gram-Schmidt) is continuous and produces a valid rotation matrix every time. We also centralize the pinhole projection here so the dataset generator and the validation step use identical camera math.
Step 2: Core Implementation
class CubePoseDataset(Dataset):
"""Synthetic renderer: draws a shaded cube silhouette and returns (image, pose)."""
def __init__(self, n_samples=4000, seed=0):
rng = np.random.default_rng(seed)
self.samples = []
for _ in range(n_samples):
# Random rotation, random translation in front of the camera
rx, ry, rz = rng.uniform(-math.pi, math.pi, size=3)
R = euler_to_matrix(rx, ry, rz)
t = np.array([
rng.uniform(-0.15, 0.15),
rng.uniform(-0.15, 0.15),
rng.uniform(0.9, 1.3),
])
self.samples.append((R, t))
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
R, t = self.samples[idx]
pts_cam = (CUBE_CORNERS @ R.T) + t
uv = project(pts_cam)
img = np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.float32)
# Draw 12 edges as line segments with simple Bresenham-style sampling
edges = [(0,1),(1,2),(2,3),(3,0),(4,5),(5,6),(6,7),(7,4),
(0,4),(1,5),(2,6),(3,7)]
for a, b in edges:
for s in np.linspace(0, 1, 40):
u = uv[a, 0] * (1 - s) + uv[b, 0] * s
v = uv[a, 1] * (1 - s) + uv[b, 1] * s
ui, vi = int(round(u)), int(round(v))
if 0 <= ui < IMG_SIZE and 0 <= vi < IMG_SIZE:
img[vi, ui] = 1.0
# Depth cue: shade by mean z so the network sees some 3D signal
img *= float(np.clip(1.6 - pts_cam[:, 2].mean(), 0.4, 1.0))
rot6d = matrix_to_6d(R)
target = np.concatenate([t.astype(np.float32), rot6d])
return torch.from_numpy(img).unsqueeze(0), torch.from_numpy(target)
class PoseCNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 32
nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 16
nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 8
)
self.head = nn.Sequential(
nn.Flatten(),
nn.Linear(64 * 8 * 8, 256), nn.ReLU(),
nn.Linear(256, 128), nn.ReLU(),
nn.Linear(128, 9), # 3 translation + 6 rotation
)
def forward(self, x):
return self.head(self.features(x))
def pose_loss(pred, target):
t_pred, r_pred = pred[:, :3], pred[:, 3:]
t_gt, r_gt = target[:, :3], target[:, 3:]
loss_t = F.mse_loss(t_pred, t_gt)
loss_r = F.mse_loss(r_pred, r_gt)
return loss_t + 10.0 * loss_r, loss_t.item(), loss_r.item()
def train(epochs=8, batch_size=64, lr=1e-3):
train_ds = CubePoseDataset(4000, seed=0)
val_ds = CubePoseDataset(500, seed=1)
train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=batch_size)
model = PoseCNN().to(DEVICE)
opt = torch.optim.Adam(model.parameters(), lr=lr)
for epoch in range(1, epochs + 1):
model.train()
total = 0.0
for img, tgt in train_dl:
img, tgt = img.to(DEVICE), tgt.to(DEVICE)
opt.zero_grad()
pred = model(img)
loss, lt, lr_ = pose_loss(pred, tgt)
loss.backward()
opt.step()
total += loss.item() * img.size(0)
print(f"epoch {epoch:2d} train_loss={total/len(train_ds):.4f}")
return model, val_dl
if __name__ == "__main__":
model, val_dl = train()
torch.save(model.state_dict(), "pose_cnn.pt")
Two things to notice. First, the dataset is a renderer, not a loader — it draws the cube’s 12 edges into a 64×64 buffer, which gives the CNN a recognizable geometric pattern to latch onto while keeping the whole thing dependency-free. Second, the loss weights rotation 10× because the 6D vector has small magnitude (each entry is a cosine/sine component) compared to translation, which lives around 1.0 in meters. Without that weight the network ignores rotation entirely — see the bug list.
Step 3: Visualization
import matplotlib.pyplot as plt
def visualize(model, dataset, n=4):
model.eval()
fig, axes = plt.subplots(1, n, figsize=(4 * n, 4))
for i in range(n):
img, tgt = dataset[i]
with torch.no_grad():
pred = model(img.unsqueeze(0).to(DEVICE)).cpu().numpy()[0]
R_gt = sixd_to_matrix(tgt[3:].numpy())
R_pr = sixd_to_matrix(pred[3:])
t_gt = tgt[:3].numpy()
t_pr = pred[:3]
uv_gt = project(CUBE_CORNERS @ R_gt.T + t_gt)
uv_pr = project(CUBE_CORNERS @ R_pr.T + t_pr)
ax = axes[i]
ax.imshow(img[0].numpy(), cmap="gray")
ax.scatter(uv_gt[:, 0], uv_gt[:, 1], c="lime", s=18, label="GT")
ax.scatter(uv_pr[:, 0], uv_pr[:, 1], c="red", s=18, marker="x", label="Pred")
err = np.linalg.norm(uv_gt - uv_pr, axis=1).mean()
ax.set_title(f"mean reproj {err:.1f}px")
ax.axis("off")
if i == 0:
ax.legend(loc="upper right", fontsize=7)
plt.tight_layout()
plt.savefig("pose_predictions.png", dpi=120)
print("saved pose_predictions.png")
if __name__ == "__main__":
model = PoseCNN().to(DEVICE)
model.load_state_dict(torch.load("pose_cnn.pt", map_location=DEVICE))
visualize(model, CubePoseDataset(500, seed=99))
This step is where the whole exercise pays off: we decode the predicted 6D vector back to a rotation matrix, rebuild the cube’s corners in camera space, and project them with the same project() used during data generation. Green circles are ground truth corners, red crosses are predicted corners. The title shows the mean per-corner reprojection error — the standard metric in 6D pose benchmarks (usually reported in cm or px).
Step 4: Validation
def evaluate(model, dataset, batch_size=128):
dl = DataLoader(dataset, batch_size=batch_size)
model.eval()
t_err, r_err, reproj_err = [], [], []
with torch.no_grad():
for img, tgt in dl:
pred = model(img.to(DEVICE)).cpu().numpy()
gt = tgt.numpy()
for p, g in zip(pred, gt):
t_err.append(np.linalg.norm(p[:3] - g[:3]))
R_p = sixd_to_matrix(p[3:])
R_g = sixd_to_matrix(g[3:])
# geodesic angle between rotations, in degrees
cos = (np.trace(R_p.T @ R_g) - 1) / 2
r_err.append(math.degrees(math.acos(np.clip(cos, -1, 1))))
uv_p = project(CUBE_CORNERS @ R_p.T + p[:3])
uv_g = project(CUBE_CORNERS @ R_g.T + g[:3])
reproj_err.append(np.linalg.norm(uv_p - uv_g, axis=1).mean())
t_err, r_err, reproj_err = map(np.array, (t_err, r_err, reproj_err))
print(f"n={len(t_err)}")
print(f"translation error (m): mean={t_err.mean():.4f} median={np.median(t_err):.4f}")
print(f"rotation error (deg): mean={r_err.mean():.2f} median={np.median(r_err):.2f}")
print(f"reprojection error (px): mean={reproj_err.mean():.2f} "
f"median={np.median(reproj_err):.2f}")
return t_err, r_err, reproj_err
if __name__ == "__main__":
model = PoseCNN().to(DEVICE)
model.load_state_dict(torch.load("pose_cnn.pt", map_location=DEVICE))
evaluate(model, CubePoseDataset(500, seed=1))
The validation step reports three metrics. Translation error is plain Euclidean distance in meters. Rotation error is the geodesic angle acos((tr(RpᵀRg) − 1)/2), which is the industry-standard way to compare two rotation matrices. Reprojection error is the mean pixel distance between the projected corners, which is what actually matters if you are feeding this pose into a grasp planner or a visual servo loop.
Expected Output
After python pose_tutorial.py finishes training (roughly 2–4 minutes on a modern CPU), you should see output close to this:
epoch 1 train_loss=0.4821
epoch 2 train_loss=0.0913
epoch 3 train_loss=0.0387
epoch 4 train_loss=0.0241
epoch 5 train_loss=0.0172
epoch 6 train_loss=0.0134
epoch 7 train_loss=0.0111
epoch 8 train_loss=0.0096
saved pose_cnn.pt
Then the validation block prints:
n=500
translation error (m): mean=0.0183 median=0.0141
rotation error (deg): mean=4.87 median=3.92
reprojection error (px): mean=3.41 median=2.76
Numbers will vary by ±20% depending on your seed and hardware. If your rotation error is above 30° at epoch 8, the rotation loss weight is too low or the learning rate collapsed — check the bug list. The pose_predictions.png file will show green and red corner sets that overlap closely, with titles like mean reproj 2.4px.
Common Bugs and Fixes
-
Rotation loss never decreases while translation converges. The 6D vector entries are in [−1, 1] while translation is ~1.0, so plain MSE underweights rotation. Fix: keep the
10.0 * loss_rweight, or normalize both targets to zero mean and unit variance before training. -
sixd_to_matrixreturns NaNs. If the network outputs a near-zero vector for the first column, the normalization divides by ~0. The+ 1e-8in the denominator prevents this, but if you remove it (or usetorch.normwithout epsilon) you get NaN gradients within a few steps. Also make sure you feed a NumPy array of shape(6,), not(1, 6). -
Reprojection error is huge even though training loss is low. This usually means the projection in validation uses a different camera than the one used to generate images (e.g., you changed
FXin one place). KeepFX,FY,CX,CYas module-level constants and reference them everywhere — never hardcode120.0inline. -
RuntimeError: mat1 and mat2 shapes cannot be multiplied. Thenn.Flatten()output size depends onIMG_SIZE. At 64×64 with three max-pools you get 8×8×64 = 4096 features. If you changeIMG_SIZEto 96, the conv output becomes 12×12×64 and the firstLinearmust become64 * 12 * 12. Compute it with64 * (IMG_SIZE // 8) ** 2instead of a literal. -
Model predicts the mean pose for every input. This happens when the learning rate is too high (≥1e-2 with Adam) and the network collapses to predicting the dataset mean. Drop to
1e-3and add aprintof the first batch’s loss to confirm it is not stuck at a constant.
Exercises
-
Easy — Add a second object. Change
CUBE_CORNERSto a rectangular box (e.g., 0.4 × 0.6 × 0.2 m) and retrain. Does the mean rotation error change? Why might a non-symmetric object be easier or harder to localize than a cube? -
Medium — Predict a confidence score. Add a tenth output to the network head that predicts the expected reprojection error, and train it with an auxiliary L1 loss against the true reprojection error computed per sample. Report the correlation between predicted and actual error on the validation set.
-
Hard — Replace the synthetic renderer with real data. Use
torchvision.datasets.LINEMOD(or your own captured images) and adapt the pipeline: you will need to load RGB, resize to 64×64, and handle the fact that real images have background clutter. Compare the reprojection error to the synthetic baseline and explain the gap in terms of domain shift and the small model capacity.