oak detector skelet
This commit is contained in:
moscovskayaliza
@@ -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()
|
||||
Reference in New Issue
Block a user