import time
import math
import cv2
import depthai as dai
import numpy as np
import torch
from ultralytics import YOLO
# ============================================================
# Settings
# ============================================================
FRAME_SIZE = (640, 400)
CAMERA_FPS = 30
MODEL_NAME = "yolo26n-pose.pt"
POSE_IMGSZ = 640
# Person detection threshold
PERSON_CONF = 0.30
# Keypoint confidence threshold
KEYPOINT_CONF = 0.35
# Valid depth range [m]
MIN_DEPTH_M = 0.30
MAX_DEPTH_M = 5.00
# Depth ROI
DEPTH_ROI_RADIUS = 4
# radius=4 -> 9 x 9 pixels
MIN_DEPTH_PIXELS = 4
# When previous Z exists, favor depths close to it
PREVIOUS_DEPTH_GATE_M = 0.40
# When center pixels are valid, use them as spatial reference
CENTER_DEPTH_GATE_M = 0.25
# ============================================================
# Temporal filter parameters
# ============================================================
# Alpha-beta filter
FILTER_ALPHA = 0.55
FILTER_BETA = 0.08
# Spatial/temporal outlier gate
MAX_JOINT_SPEED_MPS = 6.0
BASE_POSITION_GATE_M = 0.10
# Predict this many missing frames before dropping a joint
MAX_MISSED_FRAMES = 6
# Target person reset
MAX_TARGET_LOST_FRAMES = 15
# ============================================================
# Display
# ============================================================
DISPLAY_SCALE = 1.5
# True = display xyz for all 17 keypoints
SHOW_XYZ = True
# ============================================================
# COCO 17 keypoints
# ============================================================
JOINT_NAMES = [
"Nose", # 0
"LEye", # 1
"REye", # 2
"LEar", # 3
"REar", # 4
"LShoulder", # 5
"RShoulder", # 6
"LElbow", # 7
"RElbow", # 8
"LWrist", # 9
"RWrist", # 10
"LHip", # 11
"RHip", # 12
"LKnee", # 13
"RKnee", # 14
"LAnkle", # 15
"RAnkle", # 16
]
# COCO skeleton connections
SKELETON = [
(0, 1),
(0, 2),
(1, 3),
(2, 4),
(5, 6),
(5, 7),
(7, 9),
(6, 8),
(8, 10),
(5, 11),
(6, 12),
(11, 12),
(11, 13),
(13, 15),
(12, 14),
(14, 16),
]
# ============================================================
# Device for Ultralytics
# ============================================================
USE_CUDA = torch.cuda.is_available()
DEVICE = 0 if USE_CUDA else "cpu"
print("Pose inference device:", DEVICE)
if USE_CUDA:
print("GPU:", torch.cuda.get_device_name(0))
# ============================================================
# Load pose model
# ============================================================
print("Loading pose model...")
pose_model = YOLO(MODEL_NAME)
print("Pose model loaded.")
# ============================================================
# 3-D temporal tracker for one joint
# ============================================================
class JointTracker3D:
"""
Constant-velocity alpha-beta filter.
state:
position [m]
velocity [m/s]
"""
def __init__(self):
self.position = None
self.velocity = np.zeros(3, dtype=np.float64)
self.last_time = None
self.missed = 0
def reset(self):
self.position = None
self.velocity[:] = 0.0
self.last_time = None
self.missed = 0
def get_predicted_z(self):
if self.position is None:
return None
return float(self.position[2])
def update(self, measurement, now):
# ----------------------------------------------------
# No previous state
# ----------------------------------------------------
if self.position is None:
if measurement is None:
return None, False
self.position = measurement.astype(
np.float64
)
self.velocity[:] = 0.0
self.last_time = now
self.missed = 0
return self.position.copy(), True
# ----------------------------------------------------
# dt
# ----------------------------------------------------
if self.last_time is None:
dt = 1.0 / CAMERA_FPS
else:
dt = now - self.last_time
dt = np.clip(
dt,
1.0 / 120.0,
0.20
)
# ----------------------------------------------------
# Constant velocity prediction
# ----------------------------------------------------
predicted = (
self.position
+ self.velocity * dt
)
accepted = False
# ----------------------------------------------------
# Measurement available
# ----------------------------------------------------
if measurement is not None:
error = measurement - predicted
error_norm = np.linalg.norm(error)
# Allow larger motion for larger dt
position_gate = (
BASE_POSITION_GATE_M
+ MAX_JOINT_SPEED_MPS * dt
)
# ------------------------------------------------
# Accept measurement
# ------------------------------------------------
if error_norm <= position_gate:
self.position = (
predicted
+ FILTER_ALPHA * error
)
self.velocity = (
self.velocity
+ (FILTER_BETA / dt) * error
)
# Avoid unstable velocity explosion
speed = np.linalg.norm(
self.velocity
)
if speed > MAX_JOINT_SPEED_MPS:
self.velocity *= (
MAX_JOINT_SPEED_MPS
/ speed
)
self.missed = 0
accepted = True
# ------------------------------------------------
# Outlier
# ------------------------------------------------
else:
self.position = predicted
self.velocity *= 0.90
self.missed += 1
# ----------------------------------------------------
# Measurement missing
# ----------------------------------------------------
else:
self.position = predicted
self.velocity *= 0.90
self.missed += 1
self.last_time = now
# ----------------------------------------------------
# Lost for too long
# ----------------------------------------------------
if self.missed > MAX_MISSED_FRAMES:
self.reset()
return None, False
return self.position.copy(), accepted
# ============================================================
# Create 17 independent joint trackers
# ============================================================
joint_trackers = [
JointTracker3D()
for _ in range(17)
]
def reset_joint_trackers():
for tracker in joint_trackers:
tracker.reset()
# ============================================================
# Robust spatial depth estimation
# ============================================================
def robust_joint_depth(
depth_mm,
u,
v,
previous_z_m=None
):
"""
Estimate Z from an ROI around a 2-D pose keypoint.
Strategy:
1. collect valid depth around joint
2. use previous Z when available
3. otherwise use center 3x3 as spatial reference
4. reject outliers using MAD
5. return robust median
Return:
depth [m]
or None
"""
h, w = depth_mm.shape
u = int(round(u))
v = int(round(v))
if (
u < 0 or u >= w
or v < 0 or v >= h
):
return None
r = DEPTH_ROI_RADIUS
x0 = max(0, u - r)
x1 = min(w, u + r + 1)
y0 = max(0, v - r)
y1 = min(h, v + r + 1)
patch = depth_mm[
y0:y1,
x0:x1
].astype(np.float32)
# --------------------------------------------------------
# Valid range
# --------------------------------------------------------
valid = patch[
(patch > MIN_DEPTH_M * 1000.0)
&
(patch < MAX_DEPTH_M * 1000.0)
]
if len(valid) < MIN_DEPTH_PIXELS:
return None
values = valid
# ========================================================
# Temporal Z continuity
# ========================================================
if previous_z_m is not None:
previous_mm = (
previous_z_m * 1000.0
)
near_previous = values[
np.abs(
values - previous_mm
)
<
PREVIOUS_DEPTH_GATE_M * 1000.0
]
if len(near_previous) >= MIN_DEPTH_PIXELS:
values = near_previous
# ========================================================
# Spatial continuity:
# Use center 3x3 depth as reference
# ========================================================
else:
cx0 = max(0, u - 1)
cx1 = min(w, u + 2)
cy0 = max(0, v - 1)
cy1 = min(h, v + 2)
center_patch = depth_mm[
cy0:cy1,
cx0:cx1
].astype(np.float32)
center_valid = center_patch[
(center_patch > MIN_DEPTH_M * 1000.0)
&
(center_patch < MAX_DEPTH_M * 1000.0)
]
if len(center_valid) >= 2:
center_depth = np.median(
center_valid
)
near_center = values[
np.abs(
values - center_depth
)
<
CENTER_DEPTH_GATE_M * 1000.0
]
if len(near_center) >= MIN_DEPTH_PIXELS:
values = near_center
# ========================================================
# Robust MAD filtering
# ========================================================
median = np.median(values)
mad = np.median(
np.abs(values - median)
)
# Convert MAD approximately to sigma
robust_sigma = 1.4826 * mad
tolerance_mm = max(
40.0,
3.0 * robust_sigma
)
filtered = values[
np.abs(values - median)
<= tolerance_mm
]
if len(filtered) >= MIN_DEPTH_PIXELS:
values = filtered
z_mm = np.median(values)
return float(z_mm) / 1000.0
# ============================================================
# 2-D pixel + Z -> 3-D camera coordinates
# ============================================================
def pixel_to_xyz(u, v, z_m, K):
fx = K[0, 0]
fy = K[1, 1]
cx = K[0, 2]
cy = K[1, 2]
X = (u - cx) * z_m / fx
Y = (v - cy) * z_m / fy
Z = z_m
return np.array(
[X, Y, Z],
dtype=np.float64
)
# ============================================================
# 3-D -> image coordinate
# Used to visualize FILTERED / PREDICTED skeleton
# ============================================================
def xyz_to_pixel(xyz, K):
if xyz is None:
return None
X, Y, Z = xyz
if Z <= 0:
return None
fx = K[0, 0]
fy = K[1, 1]
cx = K[0, 2]
cy = K[1, 2]
u = fx * X / Z + cx
v = fy * Y / Z + cy
return (
int(round(u)),
int(round(v))
)
# ============================================================
# Select the same person between frames
#
# For this sample:
# - first frame: largest person
# - next frames: nearest box center
#
# Intended primarily for one-person motion analysis.
# ============================================================
previous_person_center = None
target_lost_frames = 0
def select_person(result):
global previous_person_center
global target_lost_frames
if (
result.boxes is None
or len(result.boxes) == 0
or result.keypoints is None
):
target_lost_frames += 1
if target_lost_frames > MAX_TARGET_LOST_FRAMES:
previous_person_center = None
reset_joint_trackers()
return None
boxes = (
result.boxes.xyxy
.detach()
.cpu()
.numpy()
)
if len(boxes) == 0:
return None
centers = np.column_stack(
(
(boxes[:, 0] + boxes[:, 2]) / 2,
(boxes[:, 1] + boxes[:, 3]) / 2
)
)
# --------------------------------------------------------
# First detection -> largest person
# --------------------------------------------------------
if previous_person_center is None:
areas = (
(boxes[:, 2] - boxes[:, 0])
*
(boxes[:, 3] - boxes[:, 1])
)
index = int(np.argmax(areas))
reset_joint_trackers()
# --------------------------------------------------------
# Maintain spatial identity
# --------------------------------------------------------
else:
distances = np.linalg.norm(
centers - previous_person_center,
axis=1
)
index = int(
np.argmin(distances)
)
previous_person_center = (
centers[index].copy()
)
target_lost_frames = 0
return index
# ============================================================
# Colors
# ============================================================
COLOR_MEASURED = (0, 255, 0)
COLOR_PREDICTED = (0, 200, 255)
COLOR_SKELETON = (255, 180, 0)
COLOR_TEXT = (255, 255, 255)
COLOR_BOX = (255, 0, 255)
# ============================================================
# DepthAI pipeline
# ============================================================
with dai.Pipeline() as pipeline:
# --------------------------------------------------------
# RGB
# --------------------------------------------------------
color = pipeline.create(
dai.node.Camera
).build(
dai.CameraBoardSocket.CAM_A,
sensorFps=CAMERA_FPS
)
# --------------------------------------------------------
# Stereo cameras
# --------------------------------------------------------
left = pipeline.create(
dai.node.Camera
).build(
dai.CameraBoardSocket.CAM_B,
sensorFps=CAMERA_FPS
)
right = pipeline.create(
dai.node.Camera
).build(
dai.CameraBoardSocket.CAM_C,
sensorFps=CAMERA_FPS
)
# --------------------------------------------------------
# Stereo depth
# --------------------------------------------------------
stereo = pipeline.create(
dai.node.StereoDepth
)
stereo.setDefaultProfilePreset(
dai.node.StereoDepth.PresetMode.DEFAULT
)
stereo.setRectifyEdgeFillColor(0)
stereo.enableDistortionCorrection(
True
)
left.requestOutput(
FRAME_SIZE
).link(
stereo.left
)
right.requestOutput(
FRAME_SIZE
).link(
stereo.right
)
# --------------------------------------------------------
# RGB-D
#
# Depth is registered to RGB.
# --------------------------------------------------------
rgbd = pipeline.create(
dai.node.RGBD
).build(
color,
stereo,
FRAME_SIZE,
CAMERA_FPS
)
rgbd_queue = (
rgbd.rgbd.createOutputQueue(
maxSize=2,
blocking=False
)
)
# --------------------------------------------------------
# Start
# --------------------------------------------------------
pipeline.start()
print()
print("OAK-D 3D Pose started")
print("Q / ESC : quit")
print()
intrinsic_printed = False
fps_value = 0.0
last_loop_time = time.perf_counter()
# ========================================================
# Main loop
# ========================================================
while pipeline.isRunning():
# ----------------------------------------------------
# Synchronized RGB-D
# ----------------------------------------------------
rgbd_data = rgbd_queue.get()
rgb_msg = rgbd_data.getRGBFrame()
depth_msg = rgbd_data.getDepthFrame()
if rgb_msg is None or depth_msg is None:
continue
frame = rgb_msg.getCvFrame()
# Registered depth
# Default RGBD depth unit = millimeter
depth_mm = depth_msg.getCvFrame()
if depth_mm.shape[:2] != frame.shape[:2]:
raise RuntimeError(
f"RGB/Depth size mismatch: "
f"RGB={frame.shape[:2]}, "
f"Depth={depth_mm.shape[:2]}"
)
h, w = frame.shape[:2]
# ====================================================
# Intrinsic parameters FOR THIS ACTUAL RGB FRAME
#
# This includes DepthAI's crop/resize transformation.
# ====================================================
transformation = (
rgb_msg.getTransformation()
)
K = np.asarray(
transformation.getIntrinsicMatrix(),
dtype=np.float64
)
if not intrinsic_printed:
print("RGB size:", w, "x", h)
print("Intrinsic matrix:")
print(K)
print(
"fx =", K[0, 0],
"fy =", K[1, 1],
"cx =", K[0, 2],
"cy =", K[1, 2]
)
intrinsic_printed = True
# ====================================================
# Ultralytics Pose
# ====================================================
inference_start = time.perf_counter()
results = pose_model.predict(
source=frame,
imgsz=POSE_IMGSZ,
conf=PERSON_CONF,
device=DEVICE,
quantize=16 if USE_CUDA else None,
verbose=False
)
inference_ms = (
time.perf_counter()
- inference_start
) * 1000.0
result = results[0]
person_index = select_person(
result
)
now = time.perf_counter()
current_conf = np.zeros(
17,
dtype=np.float32
)
measurements = [
None
for _ in range(17)
]
# ====================================================
# Person detected
# ====================================================
if person_index is not None:
keypoints_xy = (
result.keypoints.xy[
person_index
]
.detach()
.cpu()
.numpy()
)
kp_conf_tensor = (
result.keypoints.conf
)
if kp_conf_tensor is not None:
keypoints_conf = (
kp_conf_tensor[
person_index
]
.detach()
.cpu()
.numpy()
)
else:
keypoints_conf = np.ones(
17,
dtype=np.float32
)
current_conf[:] = (
keypoints_conf
)
# ------------------------------------------------
# Draw person bounding box
# ------------------------------------------------
if result.boxes is not None:
box = (
result.boxes.xyxy[
person_index
]
.detach()
.cpu()
.numpy()
.astype(int)
)
cv2.rectangle(
frame,
(box[0], box[1]),
(box[2], box[3]),
COLOR_BOX,
1
)
# =================================================
# 17 joints:
#
# 2D pose -> robust depth -> raw 3D
# =================================================
for joint_id in range(17):
conf = float(
keypoints_conf[joint_id]
)
if conf < KEYPOINT_CONF:
continue
u = float(
keypoints_xy[joint_id, 0]
)
v = float(
keypoints_xy[joint_id, 1]
)
if (
u < 0 or u >= w
or v < 0 or v >= h
):
continue
previous_z = (
joint_trackers[
joint_id
].get_predicted_z()
)
# --------------------------------------------
# Spatially robust Depth
# --------------------------------------------
z_m = robust_joint_depth(
depth_mm,
u,
v,
previous_z_m=previous_z
)
if z_m is None:
continue
# --------------------------------------------
# 3-D measurement
# --------------------------------------------
xyz = pixel_to_xyz(
u,
v,
z_m,
K
)
measurements[
joint_id
] = xyz
# ====================================================
# Temporal filtering
# ====================================================
filtered_xyz = [
None
for _ in range(17)
]
measurement_accepted = [
False
for _ in range(17)
]
for joint_id in range(17):
xyz, accepted = (
joint_trackers[
joint_id
].update(
measurements[
joint_id
],
now
)
)
filtered_xyz[
joint_id
] = xyz
measurement_accepted[
joint_id
] = accepted
# ====================================================
# Project FILTERED 3-D skeleton onto RGB
# ====================================================
display_uv = [
xyz_to_pixel(
filtered_xyz[i],
K
)
for i in range(17)
]
# ====================================================
# Skeleton
# ====================================================
for a, b in SKELETON:
pa = display_uv[a]
pb = display_uv[b]
if pa is None or pb is None:
continue
if not (
0 <= pa[0] < w
and 0 <= pa[1] < h
and 0 <= pb[0] < w
and 0 <= pb[1] < h
):
continue
cv2.line(
frame,
pa,
pb,
COLOR_SKELETON,
2,
cv2.LINE_AA
)
# ====================================================
# Joint points + XYZ coordinates
# ====================================================
for joint_id in range(17):
xyz = filtered_xyz[
joint_id
]
uv = display_uv[
joint_id
]
if xyz is None or uv is None:
continue
u, v = uv
if not (
0 <= u < w
and 0 <= v < h
):
continue
# ------------------------------------------------
# Green:
# current measured data accepted
#
# Yellow:
# temporally predicted / held joint
# ------------------------------------------------
if measurement_accepted[
joint_id
]:
color_joint = (
COLOR_MEASURED
)
else:
color_joint = (
COLOR_PREDICTED
)
cv2.circle(
frame,
(u, v),
4,
color_joint,
-1,
cv2.LINE_AA
)
# ------------------------------------------------
# XYZ text
# ------------------------------------------------
if SHOW_XYZ:
X, Y, Z = xyz
text = (
f"{joint_id}:{JOINT_NAMES[joint_id]} "
f"({X:+.2f},{Y:+.2f},{Z:.2f})"
)
# Alternate label positions slightly
if joint_id % 2 == 0:
dy = -7
else:
dy = 12
tx = min(
max(u + 5, 0),
w - 230
)
ty = min(
max(v + dy, 10),
h - 5
)
cv2.putText(
frame,
text,
(tx, ty),
cv2.FONT_HERSHEY_SIMPLEX,
0.30,
COLOR_TEXT,
1,
cv2.LINE_AA
)
# ====================================================
# FPS
# ====================================================
loop_now = time.perf_counter()
dt_loop = (
loop_now - last_loop_time
)
last_loop_time = loop_now
if dt_loop > 0:
instant_fps = 1.0 / dt_loop
if fps_value == 0:
fps_value = instant_fps
else:
fps_value = (
0.9 * fps_value
+ 0.1 * instant_fps
)
# ====================================================
# Status
# ====================================================
status = (
f"FPS {fps_value:.1f} "
f"Pose {inference_ms:.1f} ms "
f"XYZ[m]: X=right Y=down Z=forward"
)
cv2.rectangle(
frame,
(0, 0),
(w, 24),
(0, 0, 0),
-1
)
cv2.putText(
frame,
status,
(8, 17),
cv2.FONT_HERSHEY_SIMPLEX,
0.42,
(255, 255, 255),
1,
cv2.LINE_AA
)
# ====================================================
# Display
# ====================================================
if DISPLAY_SCALE != 1.0:
display_frame = cv2.resize(
frame,
None,
fx=DISPLAY_SCALE,
fy=DISPLAY_SCALE,
interpolation=cv2.INTER_LINEAR
)
else:
display_frame = frame
cv2.imshow(
"OAK-D 3D Pose",
display_frame
)
key = cv2.waitKey(1) & 0xFF
if key == ord("q") or key == 27:
break
pipeline.stop()
cv2.destroyAllWindows()
