from __future__ import annotations
import numpy as np
JOINT_NAMES = [
"Nose",
"LEye",
"REye",
"LEar",
"REar",
"LShoulder",
"RShoulder",
"LElbow",
"RElbow",
"LWrist",
"RWrist",
"LHip",
"RHip",
"LKnee",
"RKnee",
"LAnkle",
"RAnkle",
]
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),
]
# ============================================================
# 3D alpha-beta tracker
# ============================================================
class JointTracker3D:
def __init__(
self,
nominal_fps=30.0,
alpha=0.55,
beta=0.08,
max_speed_mps=6.0,
base_gate_m=0.10,
max_missed_frames=6,
):
self.nominal_fps = (
nominal_fps
)
self.alpha = alpha
self.beta = beta
self.max_speed_mps = (
max_speed_mps
)
self.base_gate_m = (
base_gate_m
)
self.max_missed_frames = (
max_missed_frames
)
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
):
# First observation
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
/ self.nominal_fps
)
else:
dt = (
now
- self.last_time
)
dt = float(
np.clip(
dt,
1.0 / 120.0,
0.20
)
)
# Constant velocity prediction
predicted = (
self.position
+ self.velocity * dt
)
accepted = False
if measurement is not None:
error = (
measurement
- predicted
)
error_norm = float(
np.linalg.norm(error)
)
gate = (
self.base_gate_m
+ self.max_speed_mps * dt
)
# Accept measurement
if error_norm <= gate:
self.position = (
predicted
+ self.alpha * error
)
self.velocity = (
self.velocity
+ (self.beta / dt)
* error
)
speed = float(
np.linalg.norm(
self.velocity
)
)
if (
speed
> self.max_speed_mps
):
self.velocity *= (
self.max_speed_mps
/ speed
)
self.missed = 0
accepted = True
# Reject sudden jump
else:
self.position = (
predicted
)
self.velocity *= 0.90
self.missed += 1
# Pose/Depth missing
else:
self.position = predicted
self.velocity *= 0.90
self.missed += 1
self.last_time = now
if (
self.missed
> self.max_missed_frames
):
self.reset()
return None, False
return (
self.position.copy(),
accepted
)
# ============================================================
# Robust Depth sampling
# ============================================================
def robust_joint_depth(
depth_mm,
u,
v,
previous_z_m,
*,
min_depth_m=0.30,
max_depth_m=5.00,
roi_radius=4,
min_pixels=4,
previous_depth_gate_m=0.40,
center_depth_gate_m=0.25,
):
h, w = depth_mm.shape
u_i = int(round(u))
v_i = int(round(v))
if not (
0 <= u_i < w
and 0 <= v_i < h
):
return None
x0 = max(
0,
u_i - roi_radius
)
x1 = min(
w,
u_i + roi_radius + 1
)
y0 = max(
0,
v_i - roi_radius
)
y1 = min(
h,
v_i + roi_radius + 1
)
patch = depth_mm[
y0:y1,
x0:x1
].astype(
np.float32
)
min_mm = (
min_depth_m * 1000.0
)
max_mm = (
max_depth_m * 1000.0
)
values = patch[
(patch > min_mm)
&
(patch < max_mm)
]
if len(values) < min_pixels:
return None
# --------------------------------------------------------
# Temporal depth 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_pixels
):
values = (
near_previous
)
# --------------------------------------------------------
# Spatial continuity
# --------------------------------------------------------
else:
cx0 = max(
0,
u_i - 1
)
cx1 = min(
w,
u_i + 2
)
cy0 = max(
0,
v_i - 1
)
cy1 = min(
h,
v_i + 2
)
center = depth_mm[
cy0:cy1,
cx0:cx1
].astype(
np.float32
)
center = center[
(center > min_mm)
&
(center < max_mm)
]
if len(center) >= 2:
center_depth = float(
np.median(center)
)
near_center = values[
np.abs(
values
- center_depth
)
<
center_depth_gate_m
* 1000.0
]
if (
len(near_center)
>= min_pixels
):
values = (
near_center
)
# --------------------------------------------------------
# MAD outlier rejection
# --------------------------------------------------------
median = float(
np.median(values)
)
mad = float(
np.median(
np.abs(
values - median
)
)
)
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_pixels:
values = filtered
return (
float(
np.median(values)
)
/ 1000.0
)
# ============================================================
# Projection
# ============================================================
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
)
return np.array(
[
x,
y,
z_m
],
dtype=np.float64
)
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))
)
