|
|
# ====== 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()
|