Spaces:
Build error
Build error
File size: 3,509 Bytes
bc3ec38 |
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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import sys
import cv2
import numpy as np
from pydantic import BaseModel
import ultralytics
from ultralytics.yolo.engine.results import Results
# Define keypoint
class GetKeypoint(BaseModel):
NOSE: int = 0
LEFT_EYE: int = 1
RIGHT_EYE: int = 2
LEFT_EAR: int = 3
RIGHT_EAR: int = 4
LEFT_SHOULDER: int = 5
RIGHT_SHOULDER: int = 6
LEFT_ELBOW: int = 7
RIGHT_ELBOW: int = 8
LEFT_WRIST: int = 9
RIGHT_WRIST: int = 10
LEFT_HIP: int = 11
RIGHT_HIP: int = 12
LEFT_KNEE: int = 13
RIGHT_KNEE: int = 14
LEFT_ANKLE: int = 15
RIGHT_ANKLE: int = 16
class DetectKeypoint:
def __init__(self, yolov8_model='yolov8m-pose'):
self.yolov8_model = yolov8_model
self.get_keypoint = GetKeypoint()
self.__load_model()
def __load_model(self):
if not self.yolov8_model.split('-')[-1] == 'pose':
sys.exit('Model not yolov8 pose')
self.model = ultralytics.YOLO(model=self.yolov8_model)
# extract function keypoint
def extract_keypoint(self, keypoint: np.ndarray) -> list:
# nose
nose_x, nose_y = keypoint[self.get_keypoint.NOSE]
# eye
left_eye_x, left_eye_y = keypoint[self.get_keypoint.LEFT_EYE]
right_eye_x, right_eye_y = keypoint[self.get_keypoint.RIGHT_EYE]
# ear
left_ear_x, left_ear_y = keypoint[self.get_keypoint.LEFT_EAR]
right_ear_x, right_ear_y = keypoint[self.get_keypoint.RIGHT_EAR]
# shoulder
left_shoulder_x, left_shoulder_y = keypoint[self.get_keypoint.LEFT_SHOULDER]
right_shoulder_x, right_shoulder_y = keypoint[self.get_keypoint.RIGHT_SHOULDER]
# elbow
left_elbow_x, left_elbow_y = keypoint[self.get_keypoint.LEFT_ELBOW]
right_elbow_x, right_elbow_y = keypoint[self.get_keypoint.RIGHT_ELBOW]
# wrist
left_wrist_x, left_wrist_y = keypoint[self.get_keypoint.LEFT_WRIST]
right_wrist_x, right_wrist_y = keypoint[self.get_keypoint.RIGHT_WRIST]
# hip
left_hip_x, left_hip_y = keypoint[self.get_keypoint.LEFT_HIP]
right_hip_x, right_hip_y = keypoint[self.get_keypoint.RIGHT_HIP]
# knee
left_knee_x, left_knee_y = keypoint[self.get_keypoint.LEFT_KNEE]
right_knee_x, right_knee_y = keypoint[self.get_keypoint.RIGHT_KNEE]
# ankle
left_ankle_x, left_ankle_y = keypoint[self.get_keypoint.LEFT_ANKLE]
right_ankle_x, right_ankle_y = keypoint[self.get_keypoint.RIGHT_ANKLE]
return [
nose_x, nose_y, left_eye_x, left_eye_y, right_eye_x, right_eye_y,
left_ear_x, left_ear_y, right_ear_x, right_ear_y, left_shoulder_x, left_shoulder_y,
right_shoulder_x, right_shoulder_y, left_elbow_x, left_elbow_y, right_elbow_x, right_elbow_y,
left_wrist_x, left_wrist_y, right_wrist_x, right_wrist_y, left_hip_x, left_hip_y,
right_hip_x, right_hip_y, left_knee_x, left_knee_y, right_knee_x, right_knee_y,
left_ankle_x, left_ankle_y,right_ankle_x, right_ankle_y
]
def get_xy_keypoint(self, results: Results) -> list:
result_keypoint = results.keypoints.xyn.cpu().numpy()[0]
keypoint_data = self.extract_keypoint(result_keypoint)
return keypoint_data
def __call__(self, image: np.array) -> Results:
results = self.model.predict(image, save=False)[0]
return results
|