diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..6122efa --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +models/ +__pycache__/ +*.pt diff --git "a/DepthAI\343\201\256\343\203\211\343\202\255\343\203\245\343\203\241\343\203\263\343\203\210 - pages.switch-science.com.url" "b/DepthAI\343\201\256\343\203\211\343\202\255\343\203\245\343\203\241\343\203\263\343\203\210 - pages.switch-science.com.url" new file mode 100644 index 0000000..98369a8 --- /dev/null +++ "b/DepthAI\343\201\256\343\203\211\343\202\255\343\203\245\343\203\241\343\203\263\343\203\210 - pages.switch-science.com.url" @@ -0,0 +1,2 @@ +[InternetShortcut] +URL=https://pages.switch-science.com/OAK-D%20documents/Main-DepthAI%E2%80%99s%20Documentation.html diff --git a/README.md b/README.md new file mode 100644 index 0000000..0ded3c7 --- /dev/null +++ b/README.md @@ -0,0 +1,49 @@ +# OAK-D 3D Pose comparison + +## Install + +Existing packages: +- depthai +- opencv-python +- numpy + +Ultralytics: + python -m pip install -U ultralytics + +MediaPipe on Windows x64: + python -m pip install mediapipe==0.10.33 + +## Files + - oakd_pose3d.py + - pose_backends.py + - pose3d_utils.py + +## Run + +Ultralytics: + python oakd_pose3d.py --backend ultralytics + +MediaPipe Lite: + python oakd_pose3d.py --backend mediapipe --mediapipe-model lite + +MediaPipe Full: + python oakd_pose3d.py --backend mediapipe --mediapipe-model full + +MediaPipe Heavy: + python oakd_pose3d.py --backend mediapipe --mediapipe-model heavy + +The MediaPipe .task model is downloaded automatically to ./models on first use. + +# OAK-D Depth Image Preview +## Files +- rgb-d_preview.py +## Run +python rgb-d_preview.py + + +# OAK-D Depth Image Preview +## Files +- pointcloud_preview.py +## Run +python pointcloud_preview.py + diff --git a/oakd_pose3d.py b/oakd_pose3d.py new file mode 100644 index 0000000..33faa8e --- /dev/null +++ b/oakd_pose3d.py @@ -0,0 +1,997 @@ +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() \ No newline at end of file diff --git a/pointcloud_preview.py b/pointcloud_preview.py new file mode 100644 index 0000000..419aebe --- /dev/null +++ b/pointcloud_preview.py @@ -0,0 +1,630 @@ +import sys +import threading + +import depthai as dai +import numpy as np +import pyvista as pv + +from pyvistaqt import BackgroundPlotter +from PySide6.QtWidgets import QApplication + +from vtkmodules.util.numpy_support import vtk_to_numpy + + +# ============================================================ +# Settings +# ============================================================ + +FRAME_SIZE = (640, 400) +FPS = 30 + +MIN_DEPTH_M = 0.3 +MAX_DEPTH_M = 5.0 + +# 640x400 = 256000 points +# +# 2 -> 128000 points +# 4 -> 64000 points +DISPLAY_STRIDE = 2 + +# GUI update rate +VIEWER_INTERVAL_MS = 33 + + +# ============================================================ +# Convert DepthAI point cloud +# ============================================================ + +def convert_pointcloud(pcl_data, unit_scale): + """ + Convert DepthAI point cloud to fixed-size arrays. + + Return: + xyz : Nx3 float32 [m] + rgba : Nx4 uint8 + valid_count : int + + Viewer coordinates: + X : right + Y : forward + Z : up + """ + + raw = np.asarray( + pcl_data.getPoints(), + dtype=np.float32 + ) + + # -------------------------------------------------------- + # Fixed decimation + # + # Do NOT remove invalid points because the VTK point count + # must remain constant between frames. + # -------------------------------------------------------- + + raw = raw[::DISPLAY_STRIDE] + + points_m = raw * unit_scale + + z = points_m[:, 2] + + valid = np.isfinite(points_m).all(axis=1) + + valid &= z >= MIN_DEPTH_M + valid &= z <= MAX_DEPTH_M + + # -------------------------------------------------------- + # DepthAI: + # + # X = right + # Y = down + # Z = forward + # + # Viewer: + # + # X = right + # Y = forward + # Z = up + # -------------------------------------------------------- + + xyz = np.empty_like( + points_m, + dtype=np.float32 + ) + + xyz[:, 0] = points_m[:, 0] + xyz[:, 1] = points_m[:, 2] + xyz[:, 2] = -points_m[:, 1] + + # -------------------------------------------------------- + # RGBA color + # -------------------------------------------------------- + + rgba = np.zeros( + (len(xyz), 4), + dtype=np.uint8 + ) + + if np.any(valid): + + # Depth normalized to 0 ... 1 + t = ( + z[valid] - MIN_DEPTH_M + ) / ( + MAX_DEPTH_M - MIN_DEPTH_M + ) + + t = np.clip( + t, + 0.0, + 1.0 + ) + + # Simple near/far color + rgba[valid, 0] = ( + 255 * (1.0 - t) + ).astype(np.uint8) + + rgba[valid, 1] = 180 + + rgba[valid, 2] = ( + 255 * t + ).astype(np.uint8) + + rgba[valid, 3] = 255 + + # ---------------------------------------------------- + # Invalid points: + # + # Place them at an existing valid point and alpha=0. + # + # This avoids NaN and avoids changing the point count. + # ---------------------------------------------------- + + first_valid = np.flatnonzero(valid)[0] + + xyz[~valid] = xyz[first_valid] + + else: + + xyz[:] = 0.0 + + return ( + xyz, + rgba, + int(np.count_nonzero(valid)) + ) + + +# ============================================================ +# DepthAI +# ============================================================ + +with dai.Pipeline() as pipeline: + + # -------------------------------------------------------- + # Cameras + # -------------------------------------------------------- + + color = pipeline.create( + dai.node.Camera + ).build( + dai.CameraBoardSocket.CAM_A, + sensorFps=FPS + ) + + left = pipeline.create( + dai.node.Camera + ).build( + dai.CameraBoardSocket.CAM_B, + sensorFps=FPS + ) + + right = pipeline.create( + dai.node.Camera + ).build( + dai.CameraBoardSocket.CAM_C, + sensorFps=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, + FPS + ) + + pcl_queue = ( + rgbd.pcl.createOutputQueue( + maxSize=2, + blocking=False + ) + ) + + pipeline.start() + + print("DepthAI started") + print("Waiting for first point cloud...") + + # ======================================================== + # First frame + # ======================================================== + + first_data = pcl_queue.get() + + raw = np.asarray( + first_data.getPoints(), + dtype=np.float32 + ) + + raw_z = raw[:, 2] + + positive = ( + np.isfinite(raw_z) + & (raw_z > 0) + ) + + median_z = np.median( + raw_z[positive] + ) + + # -------------------------------------------------------- + # Detect XYZ unit once + # -------------------------------------------------------- + + if median_z > 100.0: + + unit_scale = 0.001 + unit_name = "millimeter" + + else: + + unit_scale = 1.0 + unit_name = "meter" + + print( + f"Raw Z range: " + f"{raw_z[positive].min():.3f} - " + f"{raw_z[positive].max():.3f}" + ) + + print( + f"Median Z: {median_z:.3f}" + ) + + print( + f"Detected unit: {unit_name}" + ) + + first_xyz, first_rgba, valid_count = ( + convert_pointcloud( + first_data, + unit_scale + ) + ) + + print( + f"Viewer points: {len(first_xyz)}" + ) + + print( + f"Valid points: {valid_count}" + ) + + # ======================================================== + # Shared latest-frame buffer + # + # Camera thread -> Qt GUI thread + # ======================================================== + + latest_lock = threading.Lock() + + latest_xyz = first_xyz.copy() + latest_rgba = first_rgba.copy() + + latest_frame_id = 0 + + stop_event = threading.Event() + + # ======================================================== + # Capture thread + # + # IMPORTANT: + # DepthAI acquisition is completely separated from Qt/VTK. + # ======================================================== + + def capture_loop(): + + nonlocal_vars = { + "frame_id": 0 + } + + while not stop_event.is_set(): + + try: + + pcl_data = pcl_queue.get() + + except Exception: + + break + + xyz, rgba, count = ( + convert_pointcloud( + pcl_data, + unit_scale + ) + ) + + with latest_lock: + + # In-place replacement of shared references + globals_placeholder[0] = xyz + globals_placeholder[1] = rgba + + nonlocal_vars["frame_id"] += 1 + + globals_placeholder[2] = ( + nonlocal_vars["frame_id"] + ) + + # -------------------------------------------------------- + # Python doesn't have writable nonlocal variables at module + # scope, therefore store current frame data in this list. + # -------------------------------------------------------- + + globals_placeholder = [ + latest_xyz, + latest_rgba, + latest_frame_id + ] + + capture_thread = threading.Thread( + target=capture_loop, + daemon=True + ) + + capture_thread.start() + + # ======================================================== + # Qt + # ======================================================== + + app = QApplication.instance() + + if app is None: + + app = QApplication(sys.argv) + + app.setQuitOnLastWindowClosed(True) + + # ======================================================== + # PyVistaQt BackgroundPlotter + # ======================================================== + + plotter = BackgroundPlotter( + app=app, + show=True, + window_size=(1280, 720), + title="OAK-D Point Cloud" + ) + + # -------------------------------------------------------- + # Create ONE PolyData + # -------------------------------------------------------- + + cloud = pv.PolyData( + first_xyz.copy() + ) + + cloud.point_data["rgba"] = ( + first_rgba.copy() + ) + + actor = plotter.add_points( + cloud, + scalars="rgba", + rgba=True, + style="points", + point_size=3, + lighting=False, + name="oak_pointcloud", + reset_camera=True + ) + + plotter.show_axes() + plotter.show_grid() + + plotter.add_text( + "OAK-D Real-time Point Cloud", + position="upper_left", + font_size=10 + ) + + plotter.reset_camera() + + # ======================================================== + # IMPORTANT: + # + # Get direct NumPy views into VTK memory. + # + # We will update THESE arrays, rather than assigning new + # PyVista arrays every frame. + # ======================================================== + + vtk_points_data = ( + cloud.GetPoints().GetData() + ) + + vtk_rgba_data = ( + cloud.GetPointData().GetArray( + "rgba" + ) + ) + + vtk_points_np = vtk_to_numpy( + vtk_points_data + ) + + vtk_rgba_np = vtk_to_numpy( + vtk_rgba_data + ) + + print() + print("VTK buffers") + print( + "points:", + vtk_points_np.shape, + vtk_points_np.dtype + ) + + print( + "rgba:", + vtk_rgba_np.shape, + vtk_rgba_np.dtype + ) + + # ======================================================== + # Qt GUI update callback + # ======================================================== + + displayed_frame_id = [-1] + + update_counter = [0] + + def update_view(): + + # Window already closed + if plotter._closed: + return + + # -------------------------------------------- + # Get newest frame from capture thread + # -------------------------------------------- + + with latest_lock: + + frame_id = globals_placeholder[2] + + if frame_id == displayed_frame_id[0]: + return + + xyz = globals_placeholder[0] + rgba = globals_placeholder[1] + + # Copy because capture thread may replace + # references while VTK is drawing. + xyz = xyz.copy() + rgba = rgba.copy() + + # -------------------------------------------- + # Sanity check + # -------------------------------------------- + + if xyz.shape != vtk_points_np.shape: + + print( + "Point array shape changed:", + xyz.shape, + vtk_points_np.shape + ) + + return + + # ==================================================== + # KEY POINT: + # + # Write DIRECTLY into the existing VTK arrays. + # ==================================================== + + vtk_points_np[:] = xyz + vtk_rgba_np[:] = rgba + + # ---------------------------------------------------- + # Explicitly tell VTK which buffers changed + # ---------------------------------------------------- + + vtk_points_data.Modified() + vtk_rgba_data.Modified() + + cloud.GetPoints().Modified() + cloud.GetPointData().Modified() + cloud.Modified() + + # Mapper update + actor.mapper.Update() + + # Render + plotter.render() + + displayed_frame_id[0] = frame_id + + update_counter[0] += 1 + + if update_counter[0] % 60 == 0: + + print( + f"Viewer updates: " + f"{update_counter[0]}" + ) + + # ======================================================== + # Qt QTimer + # + # BackgroundPlotter.add_callback() uses QTimer. + # ======================================================== + + plotter.add_callback( + update_view, + interval=VIEWER_INTERVAL_MS + ) + + # ======================================================== + # Closing + # ======================================================== + + def window_closed(): + + print("Viewer closing...") + + stop_event.set() + + app.quit() + + plotter.app_window.signal_close.connect( + window_closed + ) + + print() + print("Viewer started") + print("-----------------------------") + print("Left drag : rotate") + print("Mouse wheel : zoom") + print("Middle drag : pan") + print("Q / X : close") + print() + + # ======================================================== + # Qt event loop + # + # This is the ONLY GUI event loop. + # ======================================================== + + try: + + app.exec() + + except KeyboardInterrupt: + + print("Interrupted") + + finally: + + # Stop capture thread + stop_event.set() + + # Stop DepthAI first so blocking get() is released + pipeline.stop() + + capture_thread.join( + timeout=1.0 + ) + + print("DepthAI stopped") + print("Finished") \ No newline at end of file diff --git a/pose.py b/pose.py new file mode 100644 index 0000000..47f24e6 --- /dev/null +++ b/pose.py @@ -0,0 +1,1246 @@ +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() \ No newline at end of file diff --git a/pose3d_utils.py b/pose3d_utils.py new file mode 100644 index 0000000..a4888b5 --- /dev/null +++ b/pose3d_utils.py @@ -0,0 +1,547 @@ +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)) + ) \ No newline at end of file diff --git a/pose_backends.py b/pose_backends.py new file mode 100644 index 0000000..28ec162 --- /dev/null +++ b/pose_backends.py @@ -0,0 +1,551 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import time +import urllib.request + +import cv2 +import numpy as np + + +@dataclass +class PoseResult17: + """Common COCO-17 pose result returned by all backends.""" + xy: np.ndarray # (17, 2), pixel coordinates + conf: np.ndarray # (17,), confidence/visibility + bbox: np.ndarray | None # [x1, y1, x2, y2] + + +class PoseBackendBase: + name = "unknown" + + def __init__(self): + self.last_inference_ms = 0.0 + + def infer(self, frame_bgr: np.ndarray) -> PoseResult17 | None: + raise NotImplementedError + + def close(self) -> None: + pass + + +class UltralyticsPoseBackend(PoseBackendBase): + name = "Ultralytics" + + def __init__( + self, + model_name: str = "yolo26n-pose.pt", + imgsz: int = 640, + person_conf: float = 0.30, + ): + super().__init__() + + import torch + from ultralytics import YOLO + + self.model = YOLO(model_name) + self.imgsz = imgsz + self.person_conf = person_conf + + self.use_cuda = torch.cuda.is_available() + self.device = 0 if self.use_cuda else "cpu" + + print("[Pose] backend = Ultralytics") + print(f"[Pose] model = {model_name}") + print(f"[Pose] device = {self.device}") + + if self.use_cuda: + print(f"[Pose] GPU = {torch.cuda.get_device_name(0)}") + + def infer(self, frame_bgr: np.ndarray) -> PoseResult17 | None: + kwargs = dict( + source=frame_bgr, + imgsz=self.imgsz, + conf=self.person_conf, + device=self.device, + verbose=False, + ) + + # Current Ultralytics: + # quantize=16 replaces deprecated half=True. + if self.use_cuda: + kwargs["quantize"] = 16 + + t0 = time.perf_counter() + + result = self.model.predict(**kwargs)[0] + + self.last_inference_ms = ( + time.perf_counter() - t0 + ) * 1000.0 + + if ( + result.boxes is None + or result.keypoints is None + or len(result.boxes) == 0 + ): + return None + + # Single-subject assumption: + # choose the largest detected person. + boxes = ( + result.boxes.xyxy + .detach() + .cpu() + .numpy() + ) + + areas = ( + (boxes[:, 2] - boxes[:, 0]) + * (boxes[:, 3] - boxes[:, 1]) + ) + + index = int(np.argmax(areas)) + + xy = ( + result.keypoints.xy[index] + .detach() + .cpu() + .numpy() + .astype(np.float32) + ) + + conf_tensor = result.keypoints.conf + + if conf_tensor is None: + conf = np.ones( + 17, + dtype=np.float32 + ) + else: + conf = ( + conf_tensor[index] + .detach() + .cpu() + .numpy() + .astype(np.float32) + ) + + bbox = boxes[index].astype( + np.float32 + ) + + return PoseResult17( + xy=xy, + conf=conf, + bbox=bbox + ) + + +# ============================================================ +# MediaPipe Pose Landmarker 33 -> COCO 17 +# ============================================================ + +MEDIAPIPE_TO_COCO17 = [ + 0, # Nose + 2, # Left eye + 5, # Right eye + 7, # Left ear + 8, # Right ear + 11, # Left shoulder + 12, # Right shoulder + 13, # Left elbow + 14, # Right elbow + 15, # Left wrist + 16, # Right wrist + 23, # Left hip + 24, # Right hip + 25, # Left knee + 26, # Right knee + 27, # Left ankle + 28, # Right ankle +] + + +MEDIAPIPE_MODELS = { + + "lite": ( + "pose_landmarker_lite.task", + + "https://storage.googleapis.com/" + "mediapipe-models/pose_landmarker/" + "pose_landmarker_lite/float16/1/" + "pose_landmarker_lite.task", + ), + + "full": ( + "pose_landmarker_full.task", + + "https://storage.googleapis.com/" + "mediapipe-models/pose_landmarker/" + "pose_landmarker_full/float16/1/" + "pose_landmarker_full.task", + ), + + "heavy": ( + "pose_landmarker_heavy.task", + + "https://storage.googleapis.com/" + "mediapipe-models/pose_landmarker/" + "pose_landmarker_heavy/float16/1/" + "pose_landmarker_heavy.task", + ), +} + + +class MediaPipePoseBackend(PoseBackendBase): + + name = "MediaPipe" + + def __init__( + self, + model_variant: str = "lite", + detection_conf: float = 0.40, + presence_conf: float = 0.40, + tracking_conf: float = 0.40, + model_dir: str = "models", + ): + super().__init__() + + import mediapipe as mp + + if model_variant not in MEDIAPIPE_MODELS: + raise ValueError( + "model_variant must be " + "lite, full, or heavy" + ) + + self.mp = mp + self.model_variant = model_variant + self.last_timestamp_ms = -1 + + filename, url = ( + MEDIAPIPE_MODELS[ + model_variant + ] + ) + + model_dir_path = Path( + model_dir + ) + + model_dir_path.mkdir( + parents=True, + exist_ok=True + ) + + model_path = ( + model_dir_path / filename + ) + + # ---------------------------------------------------- + # Automatic model download + # ---------------------------------------------------- + + if not model_path.exists(): + + print( + "[Pose] downloading " + "MediaPipe model:" + ) + + print( + f" {url}" + ) + + urllib.request.urlretrieve( + url, + model_path + ) + + print( + f"[Pose] saved to " + f"{model_path}" + ) + + # ---------------------------------------------------- + # VIDEO mode: + # tracking information is reused between frames. + # ---------------------------------------------------- + + options = ( + mp.tasks.vision + .PoseLandmarkerOptions( + + base_options=( + mp.tasks.BaseOptions( + model_asset_path= + str(model_path) + ) + ), + + running_mode=( + mp.tasks.vision + .RunningMode.VIDEO + ), + + num_poses=1, + + min_pose_detection_confidence= + detection_conf, + + min_pose_presence_confidence= + presence_conf, + + min_tracking_confidence= + tracking_conf, + + output_segmentation_masks=False, + ) + ) + + self.landmarker = ( + mp.tasks.vision + .PoseLandmarker + .create_from_options( + options + ) + ) + + print( + "[Pose] backend = MediaPipe" + ) + + print( + f"[Pose] model = " + f"{model_variant}" + ) + + print( + "[Pose] mode = VIDEO" + ) + + @staticmethod + def _landmark_confidence(lm): + + values = [] + + visibility = getattr( + lm, + "visibility", + None + ) + + presence = getattr( + lm, + "presence", + None + ) + + if visibility is not None: + values.append( + float(visibility) + ) + + if presence is not None: + values.append( + float(presence) + ) + + if values: + return min(values) + + return 1.0 + + def infer( + self, + frame_bgr: np.ndarray + ) -> PoseResult17 | None: + + h, w = frame_bgr.shape[:2] + + # BGR -> RGB + frame_rgb = cv2.cvtColor( + frame_bgr, + cv2.COLOR_BGR2RGB + ) + + frame_rgb = np.ascontiguousarray( + frame_rgb + ) + + mp_image = self.mp.Image( + image_format= + self.mp.ImageFormat.SRGB, + data=frame_rgb, + ) + + # VIDEO mode requires strictly increasing timestamps. + timestamp_ms = int( + time.perf_counter() + * 1000.0 + ) + + timestamp_ms = max( + timestamp_ms, + self.last_timestamp_ms + 1 + ) + + self.last_timestamp_ms = ( + timestamp_ms + ) + + t0 = time.perf_counter() + + result = ( + self.landmarker + .detect_for_video( + mp_image, + timestamp_ms, + ) + ) + + self.last_inference_ms = ( + time.perf_counter() - t0 + ) * 1000.0 + + if not result.pose_landmarks: + return None + + landmarks = ( + result.pose_landmarks[0] + ) + + xy = np.zeros( + (17, 2), + dtype=np.float32 + ) + + conf = np.zeros( + 17, + dtype=np.float32 + ) + + # ---------------------------------------------------- + # MediaPipe 33 -> COCO17 + # ---------------------------------------------------- + + for coco_id, mp_id in enumerate( + MEDIAPIPE_TO_COCO17 + ): + + lm = landmarks[mp_id] + + xy[coco_id, 0] = ( + float(lm.x) * w + ) + + xy[coco_id, 1] = ( + float(lm.y) * h + ) + + conf[coco_id] = ( + self._landmark_confidence( + lm + ) + ) + + # ---------------------------------------------------- + # Approximate bbox from landmarks + # ---------------------------------------------------- + + good = conf >= 0.20 + + bbox = None + + if np.count_nonzero(good) >= 4: + + pts = xy[good] + + pad = 15.0 + + bbox = np.array( + [ + max( + 0.0, + float( + pts[:, 0].min() + - pad + ) + ), + + max( + 0.0, + float( + pts[:, 1].min() + - pad + ) + ), + + min( + float(w - 1), + float( + pts[:, 0].max() + + pad + ) + ), + + min( + float(h - 1), + float( + pts[:, 1].max() + + pad + ) + ), + ], + + dtype=np.float32, + ) + + return PoseResult17( + xy=xy, + conf=conf, + bbox=bbox + ) + + def close(self): + + self.landmarker.close() + + +# ============================================================ +# Factory +# ============================================================ + +def create_pose_backend( + backend: str, + *, + ultralytics_model: + str = "yolo26n-pose.pt", + ultralytics_imgsz: + int = 640, + person_conf: + float = 0.30, + mediapipe_model: + str = "lite", +): + + backend = backend.lower() + + if backend == "ultralytics": + + return UltralyticsPoseBackend( + model_name= + ultralytics_model, + imgsz= + ultralytics_imgsz, + person_conf= + person_conf, + ) + + if backend == "mediapipe": + + return MediaPipePoseBackend( + model_variant= + mediapipe_model, + ) + + raise ValueError( + f"Unknown backend: {backend}" + ) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..83db0f2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +depthai +opencv-python +numpy +pyvista +vtk +pyvistaqt +PySide6 +ultralytics +torch +mediapipe==0.10.33 diff --git a/rgb-d_preview.py b/rgb-d_preview.py new file mode 100644 index 0000000..df58db1 --- /dev/null +++ b/rgb-d_preview.py @@ -0,0 +1,170 @@ +import cv2 +import depthai as dai +import numpy as np + + +FRAME_SIZE = (640, 400) +FPS = 30 + +MIN_DEPTH_MM = 300 +MAX_DEPTH_MM = 5000 + + +def colorize_depth(depth): + valid = depth > 0 + + clipped = np.clip( + depth, + MIN_DEPTH_MM, + MAX_DEPTH_MM + ) + + depth8 = np.zeros(depth.shape, dtype=np.uint8) + + depth8[valid] = ( + (MAX_DEPTH_MM - clipped[valid]) + * 255.0 + / (MAX_DEPTH_MM - MIN_DEPTH_MM) + ).astype(np.uint8) + + color = cv2.applyColorMap( + depth8, + cv2.COLORMAP_TURBO + ) + + color[~valid] = 0 + + return color + + +with dai.Pipeline() as pipeline: + + # -------------------------------------------------- + # RGB camera + # OAK-D: CAM_A + # -------------------------------------------------- + color = pipeline.create(dai.node.Camera).build( + dai.CameraBoardSocket.CAM_A, + sensorFps=FPS + ) + + # -------------------------------------------------- + # Stereo cameras + # OAK-D: CAM_B = left + # CAM_C = right + # -------------------------------------------------- + left = pipeline.create(dai.node.Camera).build( + dai.CameraBoardSocket.CAM_B, + sensorFps=FPS + ) + + right = pipeline.create(dai.node.Camera).build( + dai.CameraBoardSocket.CAM_C, + sensorFps=FPS + ) + + # -------------------------------------------------- + # Stereo depth + # -------------------------------------------------- + stereo = pipeline.create(dai.node.StereoDepth) + + stereo.setDefaultProfilePreset( + dai.node.StereoDepth.PresetMode.DEFAULT + ) + + stereo.setRectifyEdgeFillColor(0) + stereo.enableDistortionCorrection(True) + + # Left / Right camera -> StereoDepth + left.requestOutput(FRAME_SIZE).link( + stereo.left + ) + + right.requestOutput(FRAME_SIZE).link( + stereo.right + ) + + # -------------------------------------------------- + # RGBD + # + # Important: + # autocreate=True is NOT used. + # -------------------------------------------------- + rgbd = pipeline.create(dai.node.RGBD).build( + color, + stereo, + FRAME_SIZE, + FPS + ) + + rgbd_queue = rgbd.rgbd.createOutputQueue( + maxSize=4, + blocking=False + ) + + pipeline.start() + + print("Pipeline started") + print("Q / ESC : quit") + + while pipeline.isRunning(): + + data = rgbd_queue.get() + + rgb_msg = data.getRGBFrame() + depth_msg = data.getDepthFrame() + + if rgb_msg is None or depth_msg is None: + continue + + rgb = rgb_msg.getCvFrame() + + # uint16, mm + depth = depth_msg.getCvFrame() + + depth_color = colorize_depth(depth) + + # -------------------------------------------------- + # Center distance + # -------------------------------------------------- + h, w = depth.shape + + cx = w // 2 + cy = h // 2 + + d = int(depth[cy, cx]) + + if d > 0: + text = f"{d} mm" + else: + text = "invalid" + + cv2.drawMarker( + rgb, + (cx, cy), + (0, 0, 255), + cv2.MARKER_CROSS, + 20, + 2 + ) + + cv2.putText( + rgb, + text, + (20, 40), + cv2.FONT_HERSHEY_SIMPLEX, + 1, + (0, 255, 0), + 2 + ) + + cv2.imshow("RGB", rgb) + cv2.imshow("Depth", depth_color) + + key = cv2.waitKey(1) & 0xff + + if key == ord("q") or key == 27: + break + + +cv2.destroyAllWindows() \ No newline at end of file diff --git a/test1.py b/test1.py new file mode 100644 index 0000000..6d9a4b7 --- /dev/null +++ b/test1.py @@ -0,0 +1,21 @@ +import depthai as dai + +with dai.Device() as device: + print("DepthAI version:", dai.__version__) + + print("\nCamera sensors:") + print(device.getCameraSensorNames()) + + print("\nStereo pairs:") + print(device.getStereoPairs()) + + print("\nAvailable stereo pairs:") + print(device.getAvailableStereoPairs()) + + calib = device.readCalibration() + eeprom = calib.getEepromData() + + print("\nEEPROM:") + print("Product name :", eeprom.productName) + print("Board name :", eeprom.boardName) + print("Board rev :", eeprom.boardRev) \ No newline at end of file