From e2e24a5e7eb09e43cdb163985ec50e4d6ad32863 Mon Sep 17 00:00:00 2001 From: moscovskayaliza Date: Tue, 23 Jun 2026 14:51:45 +0300 Subject: [PATCH] oak detector skelet --- .gitmodules | 4 + skeleton/oak_pose_detector.py | 200 +++++++++++++++++++++++++++++ submodules/OAK-HumanPoseEstimation | 1 + 3 files changed, 205 insertions(+) create mode 100644 .gitmodules create mode 100644 skeleton/oak_pose_detector.py create mode 160000 submodules/OAK-HumanPoseEstimation diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..e907d45 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,4 @@ +[submodule "submodules/OAK-HumanPoseEstimation"] + path = submodules/OAK-HumanPoseEstimation + url = https://github.com/kschlegel/OAK-HumanPoseEstimation.git + branch = main diff --git a/skeleton/oak_pose_detector.py b/skeleton/oak_pose_detector.py new file mode 100644 index 0000000..19f5478 --- /dev/null +++ b/skeleton/oak_pose_detector.py @@ -0,0 +1,200 @@ +# ====== 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() diff --git a/submodules/OAK-HumanPoseEstimation b/submodules/OAK-HumanPoseEstimation new file mode 160000 index 0000000..5c66292 --- /dev/null +++ b/submodules/OAK-HumanPoseEstimation @@ -0,0 +1 @@ +Subproject commit 5c66292a044f07b977354ca822e37ef14ea1423d