# ====== 33 точки mediapipe ====== ''' 0 - нос 1 - внутренний угол левого глаза 2 - центр левого глаза 3 - внешний угол левого глаза 4 - внутренний угол правого глаза 5 - центр правого глаза 6 - внешний угол правого глаза 7 - левое ухо 8 - правое ухо 9 - левый угол рта 10 - правый угол рта 11 - левое плечо 12- правое плечо 13 - левый локоть 14 - правый локоть 15 - левое запястье 16 - правое запястье 17 - левый мизинец 18 - правый мизинец 19 - левый указательный палец 20 - правй указательный палец 21- левый большой палец 22 -правый большой палец 23 - левое бедро 24 - правое бедро 25 - левое колено 26 - правое колено 27 - левая голень/лодыжка 28 - правая голень/лодыжка 29 - левая пятка 30 - правая пятка 31 - левый носок 32 - правый носок ''' # ====== 17 точек efficienthrnet ====== ''' 0 - нос 1 - левый глаз 2 - правый глаз 3 - левое ухо 4 - правое ухо 5 - левое плечо 6 - правое плечо 7 - левый локоть 8 - правый локоть 9 - левое запястье 10 - правое запястье 11 - левое бедро 12 - правое бедро 13 - левое колено 14 - правое колено 15 - левая лодыжка 16 - правая лодыжка ''' import sys from pathlib import Path REPO_PATH = Path(__file__).parent.parent / "submodules" / "OAK-HumanPoseEstimation" sys.path.append(str(REPO_PATH)) import cv2 import depthai as dai import numpy as np import daipipeline from daipipeline import create_pipeline, get_model_list from poseestimators import get_poseestimator class OakPoseDetector: EH_TO_MP = { 0: 0, # nose 1: 2, # left eye 2: 5, # right eye 3: 7, # left ear 4: 8, # right ear 5: 11, # left shoulder 6: 12, # right shoulder 7: 13, # left elbow 8: 14, # right elbow 9: 15, # left wrist 10: 16, # right wrist 11: 23, # left hip 12: 24, # right hip 13: 25, # left knee 14: 26, # right knee 15: 27, # left ankle 16: 28, # right ankle } MP_CONNECTIONS = [(EH_TO_MP[0], EH_TO_MP[1]), (EH_TO_MP[1], EH_TO_MP[3]), (EH_TO_MP[0], EH_TO_MP[2]), (EH_TO_MP[2], EH_TO_MP[4]), (EH_TO_MP[5], EH_TO_MP[6]), (EH_TO_MP[5], EH_TO_MP[7]), (EH_TO_MP[7], EH_TO_MP[9]), (EH_TO_MP[6], EH_TO_MP[8]), (EH_TO_MP[8], EH_TO_MP[10]), (EH_TO_MP[5], EH_TO_MP[11]), (EH_TO_MP[6], EH_TO_MP[12]), (EH_TO_MP[11], EH_TO_MP[12]), (EH_TO_MP[11], EH_TO_MP[13]), (EH_TO_MP[13], EH_TO_MP[15]), (EH_TO_MP[12], EH_TO_MP[14]), (EH_TO_MP[14], EH_TO_MP[16])] def __init__(self, model_type='efficienthrnet1', shaves=6, detection_threshold=0.2, **kwargs): daipipeline.MODELS_FOLDER = str(REPO_PATH / "models") self.model_type = model_type self.model_list = get_model_list() if model_type not in self.model_list: raise ValueError(f"Unknown model: {model_type}") self.model_config = self.model_list[model_type] self.kwargs = kwargs self.kwargs['shaves'] = shaves self.kwargs['detection_threshold'] = int(detection_threshold * 100) if 'decoder' not in self.kwargs: self.kwargs['decoder'] = None self.pose_estimator = get_poseestimator(self.model_config, **self.kwargs) self.pipeline = create_pipeline(self.model_config, camera=True, passthrough=False, **self.kwargs) self.device = dai.Device(self.pipeline) self.preview_queue = self.device.getOutputQueue("preview", maxSize=1, blocking=False) self.pose_queue = self.device.getOutputQueue("pose", maxSize=1, blocking=False) def get_frame_and_pose(self): preview = self.preview_queue.tryGet() if preview is None: return None, None frame = preview.getCvFrame() raw_output = self.pose_queue.tryGet() if raw_output is None: return frame, None # получаем результат в формате модели personwise_keypoints = self.pose_estimator.get_pose_data(raw_output) if personwise_keypoints.shape[0] == 0: return frame, None kps = personwise_keypoints[0] # (17, 3) – (x, y, confidence) mp_landmarks = np.zeros((33, 4), dtype=np.float32) for eh_idx, mp_idx in self.EH_TO_MP.items(): x, y, conf = kps[eh_idx] if conf > 1.0: conf = conf / 100.0 mp_landmarks[mp_idx] = [x, y, 0.0, conf] return frame, mp_landmarks def draw_landmarks(self, image, landmarks, conf_threshold=0.1): if landmarks is None: return image.copy() vis = image.copy() for (i, j) in self.MP_CONNECTIONS: if i < len(landmarks) and j < len(landmarks): if landmarks[i][3] > conf_threshold and landmarks[j][3] > conf_threshold: pt1 = (int(landmarks[i][0]), int(landmarks[i][1])) pt2 = (int(landmarks[j][0]), int(landmarks[j][1])) cv2.line(vis, pt1, pt2, (0, 255, 0), 2) for i, (x, y, z, conf) in enumerate(landmarks): if conf > conf_threshold: cv2.circle(vis, (int(x), int(y)), 3, (255, 0, 0), -1) return vis def release(self): if hasattr(self, 'device'): self.device.close() def main(): detector = OakPoseDetector(model_type='efficienthrnet1', detection_threshold=0.2) print("Запуск. Нажмите 'q' для выхода.") while True: frame, landmarks = detector.get_frame_and_pose() if frame is None: continue vis = detector.draw_landmarks(frame, landmarks) if landmarks is not None else frame cv2.imshow("Pose Estimation", vis) if cv2.waitKey(1) == ord('q'): break detector.release() cv2.destroyAllWindows() if __name__ == "__main__": main()