Facial 3D Pose Estimation

The following article is inspired from Img2Pose paper.

This is the start of my 3D journey. It will be a long one !

Euler angles Yaw, Pitch & Role. It tells how a face turns in 3D, i.e can rotate in only three fundamental ways. Left-Right (Yaw), Up-Down(Pitch) & Tilt Sideways(Roll). More info at the article.

There are a total of 68 primary facial keypoints. Once these keypoints are detected, the 3D pose of a face — namely pitch, yaw, and roll — can be estimated, allowing us to determine the face’s orientation in three-dimensional space. However, detecting and regressing these keypoints requires predicting 136 values (x, y coordinates for each keypoint), which can be challenging for most models.

As a result, an alternative indirect approach is often used: Instead of predicting all keypoints explicitly, our model predicts 9 values: 6 rotation parameters and 3 translation parameters. Using just these 9 values we can estimate key-points and pose of the face — making the learning problem significantly more tractable using even a smaller model.

From 3D World Coordinates to the 2D Image Plane

Whenever we capture a photograph of a real-world object, the camera performs a geometric transformation that maps points from the three-dimensional (3D) world onto a two-dimensional (2D) image plane. This transformation is governed by the principles of projective geometry and is commonly referred to as the camera projection model.

Mathematically, this process can be expressed as:

1. The matrix [R∣t] represents the extrinsic camera parameters, which define the rotation and translation between the world coordinate system and the camera coordinate system. They describe the pose (position and orientation) of the camera relative to the world and are independent of any specific object in the scene.
2. Intrinsic Parameters: Camera Matrix K encodes the internal characteristics of the camera and maps points from the camera coordinate system to the image plane.

In practice, the exact camera intrinsic matrix K is often unknown. Estimating K requires explicit camera calibration using known calibration patterns, which is rarely feasible in unconstrained, in-the-wild image settings. As a result, many vision systems rely on approximate intrinsics derived from image properties, most commonly the image width and height.

When calibration data is unavailable, it is common to assume that the principal point lies at the center of the image:

where W and H denote the image width and height in pixels, respectively.
This assumption holds reasonably well for most consumer cameras and introduces minimal error for many applications.

Same goes for the focal length cannot be directly recovered from a single image without additional constraints. However, a widely used heuristic is to approximate it as a function of the image size:

Under these assumptions, the intrinsic matrix becomes:

This approximation allows 3D-to-2D projection and pose estimation to proceed without explicit camera calibration.

Rotation and Translation

Let’s now focus on the two remaining unknowns: the rotation R and translation t. These parameters are not known beforehand and must be predicted by the neural network. The camera intrinsic matrix K, on the other hand, is assumed to be fixed and is defined using simple image properties.

Translation shifts the object’s position relative to the camera, affecting where it appears on the image plane. Rotation changes the object’s orientation in 3D space, determining how its shape and pose are perceived in the 2D image.

A Rotation matrix

is a mathematical tool used to turn a vector (or an object) around a fixed point (the origin) without changing its size or shape.

In the real 3D world (like in a video game engine), an object is usually rotated on all three axes. We call these Euler Angles (Roll, Pitch, Yaw).

To get the final position, we multiply the matrices. Order matters!
If you want to rotate around X, then Y, then Z, the math looks like this (remember, we read right-to-left in matrix math):

Rotation matrices are a specific subset of matrices called Special Orthogonal matrices, denoted as SO(n). They have very specific and useful properties :

1. The columns of a rotation matrix are orthogonal unit vectors.
   Unit : The length of every column is 1. (Rotations do not stretch space).
   Orthogonal : The dot product of any two columns is 0. (The axes remain perpendicular; the space does not shear or skew).
2. The Inverse is the Transpose.
3. The Determinant is +1. The determinant measures how the area (or volume) changes. Since the determinant is 1, the area of the object stays the same.

A transaltion vector

where:

• tx :​ represents horizontal displacement (left or right). Increasing moves the object to the right in the image; decreasing it moves the object to the left.
• ty :​ represents vertical displacement (up or down). Increasing moves the object downward in the image; decreasing it moves the object upward (depending on the image coordinate convention).
• tz : represents depth (distance from the camera). It controls how close or far the object is from the camera. Larger tz means object appears smaller, and vice-versa.

An important practical detail is that tz is scale-ambiguous when camera intrinsics are approximated. Increasing the focal length while proportionally increasing tz produces almost the same projection.

This is why many pose-estimation systems:

• fix K
• learn tz relative to a canonical object size (e.g., average face size).

This is a full 6DoF head pose:

• 3 DoF rotation
• 3 DoF translation

Main Idea, Approach and Dataset

Our main goal is given 2D image predict [R t]. We already have images (wider-face or similar dataset), but we don’t have target values R & t.

Coming to our main equation, i explained all the components.

In the WIDER Face dataset, each image is annotated with five facial key-points: the centers of the two eyes, the nose tip, and the two corners of the mouth.
These annotated key-points correspond to 2D image coordinates and form the left-hand side of the 3D–2D projection equation.

The camera intrinsic matrix K is fixed at the beginning of the pipeline. Since the original images vary in resolution, each detected face crop is resized to a fixed resolution of 128×128 pixels. The intrinsic parameters — focal length f and principal point (cx,cy) — are then defined accordingly, with the principal point assumed to lie at the image center and the focal length chosen proportional to the image size.

To define the corresponding 3D points (x,y,z), we rely on the mean Basel Face Model (BFM).

The BFM is a statistical 3D face model built from a large collection of high-quality 3D face scans. Its mean shape represents an average human face and contains over 53,000 densely sampled 3D facial vertices, providing a consistent and anatomically meaningful facial geometry.

We select a small subset of 3D points from the mean BFM that correspond semantically to the annotated 2D landmarks — namely, the two eye centers, the nose tip, and the mouth corners. These selected 3D points act as canonical reference locations in object space and are paired with the observed 2D keypoints from the image.

This correspondence between fixed 3D landmarks from the mean BFM model and observed 2D facial key-points allows the system to estimate the unknown rotation and translation parameters, effectively recovering the 3D pose of the face from sparse 2D annotations.

************************************************************************

To reiterate, we work with thousands of face images, each annotated with its own set of five 2D facial keypoints. These 2D locations naturally vary from image to image due to differences in identity, pose, expression, and camera viewpoint. In our formulation, the camera intrinsic matrix K is fixed, and the corresponding 3D points are also fixed: we use the same five canonical landmarks extracted from the mean BFM 3D face model for every image. With these components in place, the goal is to estimate the rotation R and translation t for each face.

At this point, a natural question arises: if each face has different 2D key-points, shouldn’t the underlying 3D points also be different for each individual? After all, no two faces are exactly the same.

The key insight is that we are not trying to recover the exact 3D facial geometry of each person. Instead, we assume that all faces can be approximated by a common canonical shape — the mean BFM face — and that the observed differences in 2D keypoints are primarily explained by changes in pose rather than identity. Under this assumption, the role of R and t is to align the fixed 3D reference points to the observed 2D keypoints in the image.

In other words, the same set of 3D landmarks is reused across all images, while the pose parameters vary per image. This simplification allows us to decouple pose estimation from identity-specific shape variation, making the problem well-posed even with sparse 2D annotations. Although this approximation ignores fine-grained facial differences, it is sufficient for robust head pose estimation, where the primary objective is to recover orientation and position rather than detailed facial geometry.

Data Preparation

import scipy.io as sio
import numpy as np
import cv2
import matplotlib.pyplot as plt
import os
import random
from pathlib import Path
import matplotlib.patches as patches
from PIL import Image
from tqdm import tqdm

import torch
from torch.utils.data import Dataset
from torchvision import transforms

from torch.optim.lr_scheduler import _LRScheduler
import math

from vit import Encoder

device = 'mps'

# https://www.kaggle.com/datasets/linhlthk19/basel-face-model
bfm = sio.loadmat("01_MorphableModel.mat")

# Mean face (3N, 1) → (N, 3)
mean_shape = bfm["shapeMU"].reshape(-1, 3)
print(mean_shape.shape)  # (~53490, 3)

# https://github.com/anilbas/BFMLandmarks/blob/master/Landmarks68_BFM.anl

lm68 = np.array([
22143,22813,22840,23250,44124,45884,47085,47668,48188,48708,
49299,50498,52457,32022,32386,32359,32979,38886,39636,40030,
40238,40433,41172,41368,41578,42011,42646,8291,8305,8314,
8320,6783,7687,8331,8977,9879,1832,3760,5050,6087,4546,
3516,10731,11758,12919,14859,13191,12157,5523,6155,7442,
8345,9506,10799,11199,10179,9277,8374,7471,6566,5909,
7322,8354,9386,10941,9141,8367,7194
]) - 1 

face_68 = mean_shape[lm68]

# normalise by nose
face_68_sub = face_68 - face_68[30] 
face_68_std = np.std(face_68 - face_68[30])
face_68 = face_68_sub / face_68_std 

# 3D landmarks from BFM mean shape
nose_3d = face_68[30]
left_eye_3d = face_68[36:42].mean(axis=0)
right_eye_3d = face_68[42:48].mean(axis=0)
left_mouth_3d = face_68[48]
right_mouth_3d = face_68[54]

FACE_3D = np.stack([
    left_eye_3d,
    right_eye_3d,
    nose_3d,
    left_mouth_3d,
    right_mouth_3d
])  # (5, 3)

From the above we get our 5 set of 3D key-points. These are fixed.

Get Saptarshi MT’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

Moving on to 2D key-points we get it from wider face dataset.

# https://github.com/deepinsight/insightface/tree/master/detection/retinaface

def parse_wider_face_landmark_txt(txt_path):
    samples = []
    current_image = None

    with open(txt_path, 'r') as f:
        for line in f:
            line = line.strip()

            if line.startswith('#'):
                # New image
                if current_image is not None:
                    samples.append(current_image)

                image_path = line[2:]  # remove "# "
                current_image = {
                    "image_path": image_path,
                    "faces": []
                }

            else:
                nums = list(map(float, line.split()))
                assert len(nums) == 20, f"Expected 20 values, got {len(nums)}"

                x, y, w, h = map(int, nums[:4])

                lm = nums[4:19]
                landmarks = []
                for i in range(0, 15, 3):
                    lx, ly, v = lm[i:i+3]
                    landmarks.append((lx, ly, int(v)))

                score = nums[19]

                face = {
                    "bbox": [x, y, w, h],
                    "landmarks": landmarks,
                    "score": score
                }

                current_image["faces"].append(face)

        if current_image is not None:
            samples.append(current_image)

    return samples

samples = parse_wider_face_landmark_txt('/Users/saptarshimallikthakur/Downloads/retinaface_gt_v1.1/train/label.txt')

def visualize_sample_plt(sample, image_root, max_faces=None):
    img_path = os.path.join(image_root, sample["image_path"])
    img = Image.open(img_path).convert("RGB")

    fig, ax = plt.subplots(1, figsize=(10, 10))
    ax.imshow(img)
    ax.axis('off')

    faces = sample["faces"]
    if max_faces is not None:
        faces = faces[:max_faces]

    for face in faces:
        x, y, w, h = face["bbox"]

        # Draw bounding box
        rect = patches.Rectangle(
            (x, y), w, h,
            linewidth=2,
            edgecolor='lime',
            facecolor='none'
        )
        ax.add_patch(rect)

        # Draw landmarks
        for (lx, ly, v), color, name in zip(
            face["landmarks"],
            LANDMARK_COLORS,
            LANDMARK_NAMES
        ):
            if v != -1:
                ax.scatter(lx, ly, c=color, s=30)
                ax.text(lx + 2, ly + 2, name, fontsize=8, color=color)

    plt.show()

LANDMARK_COLORS = [
    'blue',     # left eye
    'cyan',     # right eye
    'red',      # nose
    'yellow',   # left mouth
    'magenta'   # right mouth
]

LANDMARK_NAMES = [
    'left_eye', 'right_eye', 'nose', 'left_mouth', 'right_mouth'
]

samples = parse_wider_face_landmark_txt('/Users/saptarshimallikthakur/Downloads/retinaface_gt_v1.1/train/label.txt')

folder = '/Users/saptarshimallikthakur/Pictures/WIDER_train/images'

visualize_sample_plt(samples[100], folder)

example from the wider face dataset.

Now we have all the components. We move on to estimate [R t].

dataset = []

MIN_SCORE = 0.5
MIN_VALID_LMS = 4
MAX_REPROJ_ERR = 5# pixels

CROP_SIZE = 128
f         = CROP_SIZE * 2  # w + h or you can set f = w or h
cx  = cy  = CROP_SIZE / 2  # w/2 and h/2

K = np.array([
    [f, 0, cx],
    [0, f, cy],
    [0, 0, 1]
], dtype=np.float32)

def bbox_xywh_to_xyxy(bbox):
    x, y, w, h = bbox
    return x, y, x + w, y + h

def landmarks_to_crop(landmarks, bbox, crop_size):
    x1, y1, x2, y2 = bbox_xywh_to_xyxy(bbox)

    face_2d = []
    for lx, ly, v in landmarks:
        cx = (lx - x1) / (x2 - x1)
        cy = (ly - y1) / (y2 - y1)
        face_2d.append([cx * crop_size, cy * crop_size])

    return np.array(face_2d, dtype=np.float32)


def estimate_pose(face_3d, face_2d, K):
    face_3d = face_3d.astype(np.float32)
    face_2d = face_2d.astype(np.float32)

    success, rvec, tvec = cv2.solvePnP(
        face_3d,
        face_2d,
        K,
        None,
        flags=cv2.SOLVEPNP_EPNP
    )

    if success:
        success, rvec, tvec = cv2.solvePnP(
            face_3d,
            face_2d,
            K,
            None,
            rvec,
            tvec,
            useExtrinsicGuess=True,
            flags=cv2.SOLVEPNP_ITERATIVE
        )

    if not success:
        return None, None

    R, _ = cv2.Rodrigues(rvec)
    T = tvec.reshape(3)

    return R, T


def process_face(face, face_3d, K, crop_size, min_score=0.4):
    if face["score"] < min_score:
        return None

    face_2d = landmarks_to_crop(
        face["landmarks"],
        face["bbox"],
        crop_size
    )

    R, T = estimate_pose(face_3d, face_2d, K)
    return R, T


def landmarks_to_crop_with_vis(landmarks, bbox, crop_size):
    x, y, w, h = bbox
    x1, y1, x2, y2 = x, y, x + w, y + h

    lms_crop = []
    for lx, ly, v in landmarks:
        if v == -1:
            lms_crop.append([0.0, 0.0, 0])
        else:
            cx = (lx - x1) / (x2 - x1) * crop_size
            cy = (ly - y1) / (y2 - y1) * crop_size
            lms_crop.append([cx, cy, v])

    return np.array(lms_crop, dtype=np.float32)


def valid_landmark_count(landmarks):
    return sum(1 for _, _, v in landmarks if v != -1)

def project_3d(face_3d, R, T, K):
    pts = (R @ face_3d.T + T[:, None]).T
    pts = pts / pts[:, 2:3]
    pts = (K @ pts.T).T[:, :2]
    return pts

for sample_idx, sample in enumerate(samples):
    for face_idx, face in enumerate(sample["faces"]):

        # 1. score filter
        if face["score"] < MIN_SCORE:
            continue

        # 2. landmark availability filter
        if valid_landmark_count(face["landmarks"]) < MIN_VALID_LMS:
            continue

        # 3. convert landmarks → crop space
        face_2d_crop = landmarks_to_crop_with_vis(
            face["landmarks"],
            face["bbox"],
            CROP_SIZE
        )

        # use only (x,y) for PnP
        face_2d_xy = face_2d_crop[:, :2]

        # 4. estimate pose
        R, T = estimate_pose(FACE_3D, face_2d_xy, K)
        if R is None:
            continue

        # 5. sanity check (optional but HIGHLY recommended)
        if T[2] <= 0:
            continue  # face behind camera = garbage

        # 5. reprojection error filter (CRITICAL)
        proj = project_3d(FACE_3D, R, T, K)
        reproj_err = np.mean(np.linalg.norm(proj - face_2d_xy, axis=1))

        if reproj_err > MAX_REPROJ_ERR:
            continue

        dataset.append({
            "image_path": sample["image_path"],
            "bbox": face["bbox"],
            "landmarks_img": face["landmarks"],
            "landmarks_crop": face_2d_crop.tolist(),
            "R": R.tolist(),
            "T": T.tolist()
        })

    # progress log
    if sample_idx % 100 == 0:
        print(f"Processed {sample_idx}/{len(samples)} images, "
              f"dataset size = {len(dataset)}")

Finally our dataset preparation is complete. We now have our target values i.e [R t]. Given an image we predict R & t.

def crop_and_resize(img_path, bbox, crop_size):
    img = Image.open(img_path).convert("RGB")
    x, y, w, h = bbox
    crop = img.crop((x, y, x + w, y + h))
    crop = crop.resize((crop_size, crop_size), Image.BILINEAR)
    return np.array(crop)

def visualize_pose_sample(sample_item, image_root):
    img_path = f"{image_root}/{sample_item['image_path']}"
    bbox = sample_item["bbox"]

    crop = crop_and_resize(img_path, bbox, CROP_SIZE)

    lms_crop = np.array(sample_item["landmarks_crop"])[:, :2]
    R = np.array(sample_item["R"])
    T = np.array(sample_item["T"])

    proj = project_3d(FACE_3D, R, T, K)

    plt.figure(figsize=(4, 4))
    plt.imshow(crop)
    plt.axis("off")

    # GT landmarks
    plt.scatter(
        lms_crop[:, 0],
        lms_crop[:, 1],
        c="red",
        s=40,
        label="GT landmarks"
    )

    # Reprojected landmarks
    plt.scatter(
        proj[:, 0],
        proj[:, 1],
        c="cyan",
        s=30,
        marker="x",
        label="Reprojected"
    )

    for i, name in enumerate(
        ["L-eye", "R-eye", "Nose", "L-mouth", "R-mouth"]
    ):
        plt.text(
            lms_crop[i, 0] + 2,
            lms_crop[i, 1] + 2,
            name,
            color="red",
            fontsize=8
        )

    plt.legend()
    plt.show()

import random
for _ in range(3):
    visualize_pose_sample(random.choice(dataset), folder)

class WIDERPoseDataset(Dataset):
    def __init__(self, data, img_dir):
        self.data = data
        self.img_dir = img_dir
        
        # Standard ImageNet normalization
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((CROP_SIZE,CROP_SIZE)),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        
        # Load Image
        img_path = os.path.join(self.img_dir, item['image_path'])
        img      = cv2.imread(img_path)
        img      = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        x,y,w,h  = item['bbox']
        crop     = img[y:y+h, x:x+w]
        crop     = self.transform(crop)

        # Prepare Labels
        R_mat = np.array(item['R'], dtype=np.float32) # (3,) if Rodrigues or (3,3)
        T = np.array(item['T'], dtype=np.float32) # (3,)

        # Convert Matrix to 6D representation (First 2 columns flattened)
        # R6d = [r11, r21, r31, r12, r22, r32]
        pose_6d = R_mat[:, :2].T.flatten() 
        
        # Combine Pose [6D Rot, 3D Trans] -> 9 values
        label = np.concatenate([pose_6d, T])

        # Ground Truth 2D points for Reprojection Loss
        lms_2d = torch.tensor(item['landmarks_crop'], dtype=torch.float32) # (5, 2)
        
        return {
            "image": crop,           
            "target": torch.tensor(label, dtype=torch.float32),
            "lms_2d": lms_2d                 
        }

final_dataset = WIDERPoseDataset(dataset, folder)
dataloader = torch.utils.data.DataLoader(final_dataset, batch_size=32, shuffle=True)

del final_dataset, dataset

You might also be wondering why we took first 2 columns and not third ?

Well the answer lies in Gram-Schmidt method technique presented in 2019, in a paper by Zhou et al. (“On the Continuity of Rotation Representations in Neural Networks”). More on that in next section.

Modelling

We will be using pure ViT for feature extraction and then use small head for predicting the 9 values. Link to Encoder is in my previous article.

EMB_DIM    = 192
NUM_HEADS  = EMB_DIM // 32
PATCH_SIZE = 8
DEPTH      = 6

class VitPose(torch.nn.Module):
    def __init__(self, emb_dim=192):
        super().__init__()
        self.encoder = Encoder(patch_size=PATCH_SIZE, embed_dim=emb_dim, depth=DEPTH, num_heads=NUM_HEADS, mlp_ratio=3)

        num_patches = (CROP_SIZE // PATCH_SIZE) ** 2  
        flattened_dim = num_patches * emb_dim   

        self.norm = torch.nn.LayerNorm(flattened_dim)
        self.head = torch.nn.Sequential(
            torch.nn.Linear(flattened_dim, 32),
            torch.nn.ReLU(),
            torch.nn.Linear(32 , 9)  # 6D Rot (6) + 3D Trans (3)
        )

        # --- APPLY INITIALIZATION ---
        self.apply(self._init_weights)
        
        # Specific override for the Head (Optional but recommended for Regression)
        # Initialize the final layer to near-zero to start with a "neutral" prediction
        torch.nn.init.constant_(self.head[-1].weight, 0)
        torch.nn.init.constant_(self.head[-1].bias, 0)

    def _init_weights(self, m):
        """
        ViT weight initialization
        :param m: module
        """
        if isinstance(m, torch.nn.Linear):
            # Truncated Normal (std=0.02) is the standard for ViTs
            torch.nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                torch.nn.init.constant_(m.bias, 0)
        
        elif isinstance(m, torch.nn.LayerNorm):
            torch.nn.init.constant_(m.bias, 0)
            torch.nn.init.constant_(m.weight, 1.0)

    def forward(self, img):
        x    = self.encoder(img) # (B, N, D)
        x    = x.flatten(1)
        x    = self.norm(x)
        pose = self.head(x) # (B,9)
        return pose

model = VitPose(EMB_DIM)
model.to(device)

total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"Total params     : {total:,}")
print(f"Trainable params : {trainable:,}")

Total params     : 4,648,073
Trainable params : 4,648,073

If you are training a Neural Network to predict pose,
Predicting a raw Rotation Matrix is mathematically broken.

The Problem: The Orthogonality Constraint.
A rotation matrix has strict rules:

1. Columns must be unit length.
2. Columns must be perpendicular (orthogonal).
3. Determinant must be 1.

A Neural Network outputs unconstrained floats. It may not fulfil all conditions.

Instead we predict 6 numbers.
You ask the network to predict two 3D vectors, let’s call them ‘vector a & b’.
These represent the first two columns of a rotation matrix (the X-axis and the Y-axis).

def compute_rotation_matrix_from_6d(poses):
    """
    Constructs rotation matrix from 6D representation.
    poses: (B, 6)
    Returns: (B, 3, 3)
    """
    x_raw = poses[:, 0:3] # First column
    y_raw = poses[:, 3:6] # Second column

    # Gram-Schmidt Orthogonalization
    x = x_raw / (torch.norm(x_raw, dim=1, keepdim=True) + 1e-8)
    z = torch.cross(x, y_raw, dim=1)
    z = z / (torch.norm(z, dim=1, keepdim=True) + 1e-8)
    y = torch.cross(z, x, dim=1)

    # Stack columns
    x = x.unsqueeze(2)
    y = y.unsqueeze(2)
    z = z.unsqueeze(2)
    matrix = torch.cat((x, y, z), 2) # (B, 3, 3)
    return matrix

Loss & Training

The loss has 2 components, MSE loss on prediction and Projection loss i.e how far are the predicted-projected 2D points from original 2D landmarks.

def batch_project_3d(points_3d, R, T, K):
    """
    points_3d: (B, 3, N) or (1, 3, N)
    R:         (B, 3, 3)
    T:         (B, 3, 1)
    K:         (B, 3, 3) or (1, 3, 3)
    """
    # 1. Camera Transformation: P_cam = R * X + T
    # bmm = Batch Matrix Multiply
    # R (B,3,3) * points (B,3,N) -> (B,3,N)
    cam_points = torch.bmm(R, points_3d) + T 

    # 2. Perspective Projection: P_uv = K * P_cam
    uv_homo = torch.bmm(K, cam_points)

    # 3. Normalize: u = x/z, v = y/z
    # Avoid div by zero with 1e-8
    uv_points = uv_homo[:, :2, :] / (uv_homo[:, 2:3, :] + 1e-8)
    
    # Returns (B, 2, N)
    return uv_points

class ReprojectionLoss(torch.nn.Module):
    def __init__(self, face_3d_coords, K_matrix):
        super().__init__()
        # Register as buffers so they move to GPU automatically
        self.register_buffer('face_3d', torch.tensor(face_3d_coords, dtype=torch.float32).t().unsqueeze(0)) # (1, 3, N)
        self.register_buffer('K', torch.tensor(K_matrix, dtype=torch.float32).unsqueeze(0)) # (1, 3, 3)

    def forward(self, pred_pose, gt_proj):
        batch_size = pred_pose.shape[0]

        # Expand K and Face3D to match batch size
        # (1, 3, N) -> (B, 3, N)
        face_3d_batch = self.face_3d.repeat(batch_size, 1, 1)
        K_batch       = self.K.repeat(batch_size, 1, 1)

        # Reconstruct Matrices
        pred_R = compute_rotation_matrix_from_6d(pred_pose[:, :6])
        pred_T = pred_pose[:, 6:].unsqueeze(2) # (B, 3, 1)

        # --- USE THE FUNCTION HERE ---
        pred_proj = batch_project_3d(face_3d_batch, pred_R, pred_T, K_batch).permute(0, 2, 1) # (B, N, 2)

        # 4. FIX: Slice Ground Truth to match (x, y)
        if gt_proj.shape[-1] == 3:
            gt_proj_2d = gt_proj[:, :, :2] # Drop visibility column
        else:
            gt_proj_2d = gt_proj

        loss = torch.nn.functional.l1_loss(pred_proj, gt_proj_2d)
        return loss

proj_loss_fn = ReprojectionLoss(FACE_3D, K).to(device)
mse_loss_fn  = torch.nn.MSELoss().to(device)

We will be training for 50 epochs using cosine scheduled learning rate and Adam optimizer.

class CustomLRScheduler(_LRScheduler):
    def __init__(self, optimizer, total_steps, warmup_steps_ratio, min_lr, max_lr, last_epoch=-1):
        self.total_steps = total_steps
        self.warmup_steps = int(total_steps * warmup_steps_ratio)
        self.min_lr = min_lr
        self.max_lr = max_lr
        super(CustomLRScheduler, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        current_step = self.last_epoch
        if current_step < self.warmup_steps:
            lr = self.min_lr + (self.max_lr - self.min_lr) * (current_step / self.warmup_steps)
        else:
            progress = (current_step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
            decay_factor = math.exp(-1.0 * progress)
            lr = self.min_lr + (self.max_lr - self.min_lr) * decay_factor
            lr = max(lr, self.min_lr)
        return [lr for _ in self.base_lrs]
    
EPOCHS = 100
TOTAL_STEPS = len(dataloader)*EPOCHS

LEARNING_RATE_MIN = 1e-6
LEARNING_RATE_MAX = 1e-3
WARMUP_RATIO = 0.01

optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE_MAX, betas=(0.9, 0.98), eps=1e-9)

scheduler = CustomLRScheduler(
    optimizer,TOTAL_STEPS,
    warmup_steps_ratio=WARMUP_RATIO,
    min_lr=LEARNING_RATE_MIN,
    max_lr=LEARNING_RATE_MAX
)

model.train()

for epoch in range(EPOCHS):

    train_total_loss = train_mse_loss = train_proj_loss = 0.0

    progress_bar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{EPOCHS}", leave=True, unit='batch')

    for batch in progress_bar:

        imgs          = batch['image'].to(device)
        target        = batch["target"].to(device) 
        lms_2d        = batch['lms_2d'].to(device)                

        pred = model(imgs)                          
        mse_loss  = mse_loss_fn(pred, target)
        proj_loss = 0.01 * proj_loss_fn(pred, lms_2d)

        loss = mse_loss + proj_loss

        train_total_loss  += loss.item()
        train_mse_loss    += mse_loss.item()
        train_proj_loss   += proj_loss.item()

        optimizer.zero_grad()
        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()

        progress_bar.set_postfix(tot_loss=loss.item(), mse_loss=mse_loss.item(), proj_loss=proj_loss.item(), lr=f"{scheduler.get_last_lr()[0]:.6f}")

    loss_tensor    = torch.tensor([train_total_loss, train_mse_loss, train_proj_loss], device=device)
    avg_total_loss = loss_tensor[0].item() / len(dataloader)
    avg_mse_loss   = loss_tensor[1].item() / len(dataloader)
    avg_proj_loss  = loss_tensor[2].item() / len(dataloader)

    print(f"Epoch {epoch+1} | Total_Loss {avg_total_loss:.4f} | MSE_Loss {avg_mse_loss:.4f} | Proj_Loss {avg_proj_loss:.4f} ")
    
    torch.save(model.state_dict(), 'model.pt')

Epoch 1/100: 100%|██████████| 753/753 [04:13<00:00,  2.97batch/s, lr=0.001000, mse_loss=0.149, proj_loss=0.0801, tot_loss=0.229] 
Epoch 1 | Total_Loss 1.0628 | MSE_Loss 0.7492 | Proj_Loss 0.3136 
Epoch 2/100: 100%|██████████| 753/753 [04:22<00:00,  2.87batch/s, lr=0.000990, mse_loss=0.0748, proj_loss=0.0591, tot_loss=0.134]
Epoch 2 | Total_Loss 0.1997 | MSE_Loss 0.1246 | Proj_Loss 0.0751 
.
.
.
Epoch 49/100: 100%|██████████| 753/753 [04:11<00:00,  2.99batch/s, lr=0.000616, mse_loss=0.0145, proj_loss=0.0339, tot_loss=0.0484] 
Epoch 49 | Total_Loss 0.0400 | MSE_Loss 0.0093 | Proj_Loss 0.0307 
Epoch 50/100: 100%|██████████| 753/753 [04:11<00:00,  3.00batch/s, lr=0.000610, mse_loss=0.00589, proj_loss=0.0276, tot_loss=0.0335]
Epoch 50 | Total_Loss 0.0397 | MSE_Loss 0.0091 | Proj_Loss 0.0306 
.
.
.
Epoch 99/100: 100%|██████████| 753/753 [04:14<00:00,  2.96batch/s, lr=0.000372, mse_loss=0.00213, proj_loss=0.0213, tot_loss=0.0235]
Epoch 99 | Total_Loss 0.0297 | MSE_Loss 0.0044 | Proj_Loss 0.0252 
Epoch 100/100: 100%|██████████| 753/753 [04:14<00:00,  2.96batch/s, lr=0.000369, mse_loss=0.00306, proj_loss=0.0264, tot_loss=0.0295]
Epoch 100 | Total_Loss 0.0297 | MSE_Loss 0.0045 | Proj_Loss 0.0252 

Note: You might have to play with projection loss weight if you loss ain’t improving.

Inference

Now in inference we will just pass an image.

• The model will output 9 values which is 6 for rotation matrix corresponding to first two columns and 3 for translation vector.
• Using Gram-Schmidt formula we then compute all 9 rotation values.
• Now using the mean face BFM model (which as explained is static and pre-defined) we can project the 3D points accurately on the face image, hence indirectly solving the facial key-point detection problem efficiently !

inference_model = VitPose(EMB_DIM)
inference_model.load_state_dict(torch.load("model.pt", map_location=device))
inference_model.to(device)
inference_model.eval()

index      = 110
image_root = '/Users/saptarshimallikthakur/Pictures/WIDER_train/images'
sample     = samples[index]
img_path   = os.path.join(image_root, sample["image_path"])

# Standard ImageNet normalization
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((CROP_SIZE,CROP_SIZE)),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

FACE_3D_T = torch.tensor(FACE_3D, dtype=torch.float32).to(device).t().unsqueeze(0)
K_T       = torch.tensor(K, dtype=torch.float32).unsqueeze(0).to(device)

img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

valid_faces = [f for f in sample['faces'] if f['score'] > 0.5]

# Big vertical figure
plt.figure(figsize=(8, 6 * len(valid_faces)))

for idx, face in enumerate(valid_faces):

    x, y, w, h = face['bbox']
    crop      = img[y:y+h, x:x+w]
    crop_t    = transform(crop).unsqueeze(0).to(device)
    viz_crop  = cv2.resize(crop, (CROP_SIZE, CROP_SIZE))

    with torch.no_grad():
        pred_pose = inference_model(crop_t)
        pred_R    = compute_rotation_matrix_from_6d(pred_pose[:, :6])
        pred_T    = pred_pose[:, 6:].unsqueeze(2)

        pred_proj = batch_project_3d(FACE_3D_T, pred_R, pred_T, K_T)
        points    = pred_proj.squeeze(0).cpu().numpy()

    # One row, one plot
    ax = plt.subplot(len(valid_faces), 1, idx + 1)
    ax.imshow(viz_crop)

    ax.scatter(points[0], points[1], s=200, c='cyan', marker='.')

    ax.scatter(points[0, 2], points[1, 2], s=200, c='red', marker='x')

    # ax.set_title(f"Face {idx+1} (Score: {face['score']:.2f})", fontsize=14)
    ax.axis('off')

plt.tight_layout()
plt.show()

We can predict all 68 facial key-points too and not just 5 which we used for [R t] estimation !

index      = 103
image_root = '/Users/saptarshimallikthakur/Pictures/WIDER_train/images'
sample     = samples[index]
img_path   = os.path.join(image_root, sample["image_path"])

# Standard ImageNet normalization
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((CROP_SIZE,CROP_SIZE)),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

FACE_3D_T = torch.tensor(face_68, dtype=torch.float32).to(device).t().unsqueeze(0)
K_T       = torch.tensor(K, dtype=torch.float32).unsqueeze(0).to(device)

img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

valid_faces = [f for f in sample['faces'] if f['score'] > 0.5]

# Big vertical figure
plt.figure(figsize=(8, 6 * len(valid_faces)))

for idx, face in enumerate(valid_faces):

    x, y, w, h = face['bbox']
    crop      = img[y:y+h, x:x+w]
    crop_t    = transform(crop).unsqueeze(0).to(device)
    viz_crop  = cv2.resize(crop, (CROP_SIZE, CROP_SIZE))

    with torch.no_grad():
        pred_pose = inference_model(crop_t)
        pred_R    = compute_rotation_matrix_from_6d(pred_pose[:, :6])
        pred_T    = pred_pose[:, 6:].unsqueeze(2)

        pred_proj = batch_project_3d(FACE_3D_T, pred_R, pred_T, K_T)
        points    = pred_proj.squeeze(0).cpu().numpy()

    # One row, one plot
    ax = plt.subplot(len(valid_faces), 1, idx + 1)
    ax.imshow(viz_crop)

    ax.scatter(points[0], points[1], s=200, c='cyan', marker='.')

    # ax.scatter(points[0, 2], points[1, 2], s=200, c='red', marker='x')

    ax.set_title(f"Face {idx+1} (Score: {face['score']:.2f})", fontsize=14)
    ax.axis('off')

plt.tight_layout()
plt.show()

You can improve results by using a better pre-trained backbone (frozen) for feature extraction and just add a trainable head for prediction, unlike in my case where every component was initialised and trained from scratch.

Additional Notes

There is still one thing that I didn’t cover and that is rvec which you might have seen in :

R, _ = cv2.Rodrigues(rvec)

where R is the rotation matrix which we talked in-depth in article above, while rvec is the rotation vector of shape [1, 3].

rvec is a compact way to represent a rotation using only 3 numbers, instead of the 9 numbers required for a Rotation Matrix. It is a clever mathematical hack that packs two pieces of information into a single 3D vector (x, y, z).

1. The Direction: The direction the vector points indicates the Axis of Rotation (the imaginary line the object spins around).
2. The Magnitude (Length): The length of the vector indicates the Angle of Rotation (theta) in radians.

In computer vision (like calibrating a camera or solving solvePnP to find a face’s pose), the computer tries to “guess” the rotation.

• If the solver guesses a Matrix, it has to guess 9 numbers. But it must also ensure those 9 numbers form a valid orthogonal matrix (columns perpendicular, length 1). This is extremely difficult for an optimizer.
• If the solver guesses an rvec, it can guess any 3 numbers. There are no constraints. It creates a smooth “search space” for the algorithm to find the best rotation.

Still we don’t use it in neural network prediction.

Yaw, Pitch & Roll

As we discussed, these are the three Euler Angles that describe the 3D orientation (rotation) of an object in space relative to a camera or a coordinate system.

• Left-Right (Yaw),
• Up-Down(Pitch) &
• Tilt Sideways(Roll).

It is used in Driver Monitoring Systems (DMS):

• Yaw: Is the driver looking at the road or out the side window?
• Pitch: Is the driver nodding off (falling asleep)?

The values are computed from rotation matrix.

def compute_euler_angles_from_rotation_matrices(R):
    """
    R: (B, 3, 3) pytorch tensor or (3, 3) numpy array
    Returns: yaw, pitch, roll in degrees
    """
    # If numpy, convert to tensor for easier handling (or keep pure numpy)
    is_tensor = torch.is_tensor(R)
    
    if is_tensor:
        # Check if singular (Gimbal Lock)
        sy = torch.sqrt(R[:, 0, 0] * R[:, 0, 0] +  R[:, 1, 0] * R[:, 1, 0])
        singular = sy < 1e-6

        x = torch.zeros_like(sy)
        y = torch.zeros_like(sy)
        z = torch.zeros_like(sy)

        # Standard Case
        x[~singular] = torch.atan2(R[~singular, 2, 1], R[~singular, 2, 2])
        y[~singular] = torch.atan2(-R[~singular, 2, 0], sy[~singular])
        z[~singular] = torch.atan2(R[~singular, 1, 0], R[~singular, 0, 0])

        # Singular Case (Gimbal lock)
        x[singular] = torch.atan2(-R[singular, 1, 2], R[singular, 1, 1])
        y[singular] = torch.atan2(-R[singular, 2, 0], sy[singular])
        z[singular] = 0

        # Convert Radians to Degrees
        return x * (180 / math.pi), y * (180 / math.pi), z * (180 / math.pi)
    
    else:
        # Numpy Version
        sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])
        singular = sy < 1e-6

        if not singular:
            x = math.atan2(R[2,1] , R[2,2])
            y = math.atan2(-R[2,0], sy)
            z = math.atan2(R[1,0], R[0,0])
        else:
            x = math.atan2(-R[1,2], R[1,1])
            y = math.atan2(-R[2,0], sy)
            z = 0

        return np.array([x, y, z]) * (180 / math.pi)

index      = 103
image_root = '/Users/saptarshimallikthakur/Pictures/WIDER_train/images'
sample     = samples[index]
img_path   = os.path.join(image_root, sample["image_path"])

# Standard ImageNet normalization
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((CROP_SIZE,CROP_SIZE)),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

FACE_3D_T = torch.tensor(FACE_3D, dtype=torch.float32).to(device).t().unsqueeze(0)
K_T       = torch.tensor(K, dtype=torch.float32).unsqueeze(0).to(device)

img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

valid_faces = [f for f in sample['faces'] if f['score'] > 0.5]

# Big vertical figure
plt.figure(figsize=(8, 6 * len(valid_faces)))

for idx, face in enumerate(valid_faces):

    x, y, w, h = face['bbox']
    crop      = img[y:y+h, x:x+w]
    crop_t    = transform(crop).unsqueeze(0).to(device)
    viz_crop  = cv2.resize(crop, (CROP_SIZE, CROP_SIZE))

    with torch.no_grad():
        pred_pose = inference_model(crop_t)
        pred_R    = compute_rotation_matrix_from_6d(pred_pose[:, :6])

        pitch, yaw, roll = compute_euler_angles_from_rotation_matrices(pred_R)

    yaw_val   = yaw[0].item()
    pitch_val = pitch[0].item()
    roll_val  = roll[0].item()

    if pitch_val > 0:
        pitch_val -= 180
    else:
        pitch_val += 180

    ax = plt.subplot(len(valid_faces), 1, idx + 1)
    ax.imshow(viz_crop)

    ax.set_title(
        f"Yaw: {yaw_val:.2f} | Pitch: {pitch_val:.2f} | Roll: {roll_val:.2f}",
        fontsize=14
    )

    ax.axis('off')

plt.tight_layout()
plt.show()

Opencv and BFM have different coordinate axis direction.

• OpenCV (Camera): The Y-axis points DOWN (pixel 0,0 is top-left).
• BFM (3D Model): The Y-axis usually points UP (Cartesian coordinates).

When cv2.solvePnP tries to fit the “Y-Up” face model onto the “Y-Down” image, it finds that the best way to match the points is to rotate the face 180 degrees around the X-axis (Pitch).

Effectively, the math thinks: “To make this 3D face match the picture, I have to flip it upside down.”

In the below

• Yaw negative means head is positioned left
• Pitch negative means head is positioned facing up
• Roll negative means negative means head tilted right.

Small values — > [-5,5] means neutral (IMHO).

The End

And thats it folks. Happy learning !