51 lines
1 KiB
Python
51 lines
1 KiB
Python
from dataclasses import dataclass
|
|
|
|
# COCO-17 keypoint indices
|
|
NOSE = 0
|
|
L_EYE, R_EYE, L_EAR, R_EAR = 1, 2, 3, 4
|
|
L_SHOULDER, R_SHOULDER = 5, 6
|
|
L_ELBOW, R_ELBOW, L_WRIST, R_WRIST = 7, 8, 9, 10
|
|
L_HIP, R_HIP = 11, 12
|
|
L_KNEE, R_KNEE, L_ANKLE, R_ANKLE = 13, 14, 15, 16
|
|
CORE_JOINTS = (L_SHOULDER, R_SHOULDER, L_HIP, R_HIP)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Keypoint:
|
|
x: float
|
|
y: float
|
|
conf: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BBox:
|
|
x1: float
|
|
y1: float
|
|
x2: float
|
|
y2: float
|
|
|
|
@property
|
|
def width(self) -> float:
|
|
return self.x2 - self.x1
|
|
|
|
@property
|
|
def height(self) -> float:
|
|
return self.y2 - self.y1
|
|
|
|
@property
|
|
def aspect(self) -> float:
|
|
return self.width / self.height if self.height else 0.0
|
|
|
|
@property
|
|
def cx(self) -> float:
|
|
return (self.x1 + self.x2) / 2
|
|
|
|
@property
|
|
def cy(self) -> float:
|
|
return (self.y1 + self.y2) / 2
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Person:
|
|
bbox: BBox
|
|
keypoints: list # list[Keypoint], length 17
|