GitBucket
from __future__ import annotations
import argparse
import time
import cv2
import depthai as dai
import numpy as np
from pose_backends import (
create_pose_backend
)
from pose3d_utils import (
JOINT_NAMES,
SKELETON,
JointTracker3D,
pixel_to_xyz,
robust_joint_depth,
xyz_to_pixel,
)
# ============================================================
# Configuration
# ============================================================
FRAME_SIZE = (
640,
400
)
CAMERA_FPS = 30
DISPLAY_SCALE = 1.5
PERSON_CONF = 0.30
KEYPOINT_CONF = 0.35
MIN_DEPTH_M = 0.30
MAX_DEPTH_M = 5.00
DEPTH_ROI_RADIUS = 4
MIN_DEPTH_PIXELS = 4
PREVIOUS_DEPTH_GATE_M = 0.40
CENTER_DEPTH_GATE_M = 0.25
FILTER_ALPHA = 0.55
FILTER_BETA = 0.08
MAX_JOINT_SPEED_MPS = 6.0
BASE_POSITION_GATE_M = 0.10
MAX_MISSED_FRAMES = 6
SHOW_XYZ = True
# ============================================================
# 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
)
# ============================================================
# Command line
# ============================================================
def parse_args():
parser = argparse.ArgumentParser(
description=(
"OAK-D RGB-D "
"3D pose estimation"
)
)
parser.add_argument(
"--backend",
choices=[
"ultralytics",
"mediapipe"
],
default="ultralytics",
help=(
"Pose estimation backend"
),
)
parser.add_argument(
"--mediapipe-model",
choices=[
"lite",
"full",
"heavy"
],
default="lite",
help=(
"MediaPipe model"
),
)
parser.add_argument(
"--yolo-model",
default=(
"yolo26n-pose.pt"
),
help=(
"Ultralytics pose model"
),
)
parser.add_argument(
"--yolo-imgsz",
type=int,
default=640,
help=(
"Ultralytics input size"
),
)
return parser.parse_args()
# ============================================================
# Trackers
# ============================================================
def create_joint_trackers():
return [
JointTracker3D(
nominal_fps=
CAMERA_FPS,
alpha=
FILTER_ALPHA,
beta=
FILTER_BETA,
max_speed_mps=
MAX_JOINT_SPEED_MPS,
base_gate_m=
BASE_POSITION_GATE_M,
max_missed_frames=
MAX_MISSED_FRAMES,
)
for _ in range(17)
]
# ============================================================
# Draw skeleton
# ============================================================
def draw_pose(
frame,
filtered_xyz,
measurement_accepted,
K,
):
h, w = frame.shape[:2]
display_uv = [
xyz_to_pixel(
xyz,
K
)
for xyz
in filtered_xyz
]
# --------------------------------------------------------
# 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,
)
# --------------------------------------------------------
# Joints
# --------------------------------------------------------
for joint_id, xyz in enumerate(
filtered_xyz
):
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
if measurement_accepted[
joint_id
]:
joint_color = (
COLOR_MEASURED
)
else:
joint_color = (
COLOR_PREDICTED
)
cv2.circle(
frame,
(u, v),
4,
joint_color,
-1,
cv2.LINE_AA,
)
# ----------------------------------------------------
# XYZ text
# ----------------------------------------------------
if SHOW_XYZ:
x, y, z = xyz
text = (
f"{joint_id}:"
f"{JOINT_NAMES[joint_id]} "
f"({x:+.2f},"
f"{y:+.2f},"
f"{z:.2f})"
)
dy = (
-7
if joint_id % 2 == 0
else 12
)
tx = min(
max(
u + 5,
0
),
max(
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,
)
# ============================================================
# Main
# ============================================================
def main():
args = parse_args()
# --------------------------------------------------------
# Pose backend
# --------------------------------------------------------
pose_backend = (
create_pose_backend(
args.backend,
ultralytics_model=
args.yolo_model,
ultralytics_imgsz=
args.yolo_imgsz,
person_conf=
PERSON_CONF,
mediapipe_model=
args.mediapipe_model,
)
)
joint_trackers = (
create_joint_trackers()
)
try:
with dai.Pipeline() as pipeline:
# =================================================
# RGB camera
# =================================================
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,
)
# =================================================
# StereoDepth
# =================================================
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
)
# =================================================
# RGBD
# =================================================
rgbd = pipeline.create(
dai.node.RGBD
).build(
color,
stereo,
FRAME_SIZE,
CAMERA_FPS,
)
rgbd_queue = (
rgbd.rgbd
.createOutputQueue(
maxSize=2,
blocking=False,
)
)
pipeline.start()
print()
print(
"OAK-D 3D Pose started"
)
print(
f"Pose backend: "
f"{pose_backend.name}"
)
print(
"Q / ESC : quit"
)
print()
intrinsic_printed = False
fps_value = 0.0
last_loop_time = (
time.perf_counter()
)
frame_counter = 0
# =================================================
# 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()
)
depth_mm = (
depth_msg.getCvFrame()
)
if (
depth_mm.shape[:2]
!= frame.shape[:2]
):
raise RuntimeError(
"RGB/Depth size mismatch: "
f"RGB={frame.shape[:2]}, "
f"Depth={depth_mm.shape[:2]}"
)
h, w = frame.shape[:2]
# --------------------------------------------
# RGB intrinsic matrix
# --------------------------------------------
K = np.asarray(
rgb_msg
.getTransformation()
.getIntrinsicMatrix(),
dtype=np.float64,
)
if not intrinsic_printed:
print(
f"RGB size: "
f"{w} x {h}"
)
print(
"Intrinsic matrix:"
)
print(K)
intrinsic_printed = True
# =============================================
# Pose inference
# =============================================
pose = (
pose_backend.infer(
frame
)
)
measurements = [
None
] * 17
now = (
time.perf_counter()
)
# =============================================
# 2D Pose -> Depth -> XYZ
# =============================================
if pose is not None:
# ----------------------------------------
# Bounding box
# ----------------------------------------
if pose.bbox is not None:
box = (
pose.bbox
.astype(int)
)
cv2.rectangle(
frame,
(
box[0],
box[1]
),
(
box[2],
box[3]
),
COLOR_BOX,
1,
)
# ----------------------------------------
# 17 joints
# ----------------------------------------
for joint_id in range(
17
):
conf = float(
pose.conf[
joint_id
]
)
if (
conf
< KEYPOINT_CONF
):
continue
u = float(
pose.xy[
joint_id,
0
]
)
v = float(
pose.xy[
joint_id,
1
]
)
if not (
0 <= u < w
and
0 <= v < h
):
continue
previous_z = (
joint_trackers[
joint_id
]
.get_predicted_z()
)
# ------------------------------------
# Robust ROI depth
# ------------------------------------
z_m = robust_joint_depth(
depth_mm,
u,
v,
previous_z_m=
previous_z,
min_depth_m=
MIN_DEPTH_M,
max_depth_m=
MAX_DEPTH_M,
roi_radius=
DEPTH_ROI_RADIUS,
min_pixels=
MIN_DEPTH_PIXELS,
previous_depth_gate_m=
PREVIOUS_DEPTH_GATE_M,
center_depth_gate_m=
CENTER_DEPTH_GATE_M,
)
if z_m is None:
continue
# ------------------------------------
# Pixel -> XYZ
# ------------------------------------
measurements[
joint_id
] = pixel_to_xyz(
u,
v,
z_m,
K,
)
# =============================================
# Temporal 3D filtering
# =============================================
filtered_xyz = [
None
] * 17
measurement_accepted = [
False
] * 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
# =============================================
# Draw pose
# =============================================
draw_pose(
frame,
filtered_xyz,
measurement_accepted,
K,
)
# =============================================
# 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.0:
fps_value = (
instant_fps
)
else:
fps_value = (
0.9 * fps_value
+
0.1 * instant_fps
)
# =============================================
# Status
# =============================================
if (
args.backend
== "mediapipe"
):
backend_detail = (
args.mediapipe_model
)
else:
backend_detail = (
args.yolo_model
)
status = (
f"{pose_backend.name}/"
f"{backend_detail} "
f"FPS {fps_value:.1f} "
f"Pose "
f"{pose_backend.last_inference_ms:.1f} ms "
f"XYZ[m]: "
f"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.40,
(
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,
)
# =============================================
# Benchmark output
# =============================================
frame_counter += 1
if (
frame_counter
% 60
== 0
):
print(
f"{pose_backend.name}: "
f"pose="
f"{pose_backend.last_inference_ms:.1f} ms, "
f"total="
f"{fps_value:.1f} fps"
)
# =============================================
# Keyboard
# =============================================
key = (
cv2.waitKey(1)
& 0xFF
)
if (
key == ord("q")
or key == 27
):
break
pipeline.stop()
finally:
pose_backend.close()
cv2.destroyAllWindows()
if __name__ == "__main__":
main()