litert-community/6DRepNet-HeadPose-LiteRT

6DRepNet — Head pose estimation (LiteRT GPU)

On-device 6-DoF head pose estimation running fully on the LiteRT CompiledModel GPU delegate (no CPU fallback). 6DRepNet (ICIP 2022) regresses a continuous 6D rotation from a face crop — yaw / pitch / roll for driver-monitoring, AR, and attention. ~21 ms/frame on a Pixel 8a.

• Architecture: RepVGG-B1g2 backbone (deploy/re-parameterized) + 6D rotation head — pure CNN.
• Weights: thohemp/6DRepNet (300W-LP) · MIT.
• Size: 157 MB.

3D head-pose axes + yaw/pitch/roll on a face crop. Portrait: Unsplash (free license).

I/O

• Input: [1, 3, 224, 224] NCHW, RGB, ImageNet-normalized (a face crop; use a face detector, or a centered crop for a frontal demo).
• Output: [1, 6] — a continuous 6D rotation representation.

Host-side decode (6D → Euler)

Gram-Schmidt the 6D into a 3×3 rotation matrix, then read the Euler angles:

x = normalize(v[0:3]);  z = normalize(cross(x, v[3:6]));  y = cross(z, x)   # R = [x|y|z]
pitch = atan2(R21, R22);  yaw = atan2(-R20, sqrt(R00^2+R10^2));  roll = atan2(R10, R00)

GPU conversion

6DRepNet (deploy-mode RepVGG = plain 3×3 convs + ReLU) is a pure CNN → fully GPU-compatible (36/36 nodes on the delegate, 1 partition; device corr 0.9993, ~21 ms) with zero patches. The 6D→rotation→Euler decode runs host-side. Use the deploy weights (fused rbr_reparam), not the training-mode branches. CPU-exact vs PyTorch (corr 1.0).

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

val options = CompiledModel.Options(Accelerator.GPU)
val model = CompiledModel.create(context.assets, "6drepnet.tflite", options, null)
val inBufs = model.createInputBuffers()
val outBufs = model.createOutputBuffers()

inBufs[0].writeFloat(faceCropNCHW)       // [1,3,224,224] RGB, ImageNet-norm
model.run(inBufs, outBufs)
val v = outBufs[0].readFloat()           // [6]; Gram-Schmidt -> R -> yaw/pitch/roll (see above)

Python (LiteRT / ai-edge-litert)

import numpy as np
from ai_edge_litert.interpreter import Interpreter

it = Interpreter(model_path="6drepnet.tflite"); it.allocate_tensors()
inp, out = it.get_input_details(), it.get_output_details()
it.set_tensor(inp[0]["index"], x)        # [1,3,224,224] float32, RGB, ImageNet-norm
it.invoke()
v = it.get_tensor(out[0]["index"])[0]    # [6] -> Gram-Schmidt -> rotation matrix -> Euler

Conversion

Converted with litert-torch (build_6drepnet.py): loads the deploy-mode RepVGG weights and exports the 6D head (input face crop → 6D).

License

MIT (6DRepNet / thohemp). Trained on 300W-LP.