3D_SkeletonCapture.py

from ultralytics import YOLO
import pyrealsense2 as rs
import cv2
import numpy as np
import time
import socket
import json

UDP_IP = "127.0.0.1"   # MATLAB receiver IP (localhost)
UDP_PORT = 5004        # MATLAB listening port

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# Load YOLO11 Pose Model
model = YOLO("yolo11n-pose.pt")  # Or yolo11n-pose, yolo11s-pose, etc.

# Establishes RealSense pipeline and config
pipeline = rs.pipeline()
config = rs.config()

# Enables color and depth streams
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)

# Starts pipeline
pipeline.start(config)

try:
    while True:
        frame = pipeline.wait_for_frames()  # _____________ Gets frames
        color_frame = frame.get_color_frame()  # __________ Gets color frame
        depth_frame = frame.get_depth_frame()  # __________ Gets depth frame

        color = np.asanyarray(color_frame.get_data())  # __ Gets color image from frame
        depth = np.asanyarray(depth_frame.get_data())  # __ Gets depth image from frame

        # Run inference
        results = model(color)

        # Draw pose outputs on frame
        annotated = results[0].plot()
        if results[0].keypoints is not None:
            keypoints = results[0].keypoints.xy

            

            keypointsList = keypoints[0].tolist()

            for k in keypointsList:
                k.append(depth[round(k[1])-1,round(k[0])-1].item())

            msg = json.dumps(keypointsList[1:13])
            sock.sendto(msg.encode(), (UDP_IP, UDP_PORT))
            print(msg)
            #print(results[0].keypoints.conf[0][5])

        cv2.imshow("YOLO11-Pose RealSense", annotated)
        cv2.waitKey(20)
finally:
    pipeline.stop()