feat(edge): data models (Keypoint, BBox, Person) + COCO indices
This commit is contained in:
parent
a4f2ead211
commit
33e987d8d9
4 changed files with 84 additions and 0 deletions
0
edge/__init__.py
Normal file
0
edge/__init__.py
Normal file
51
edge/models.py
Normal file
51
edge/models.py
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
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
|
||||
12
pyproject.toml
Normal file
12
pyproject.toml
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
[project]
|
||||
name = "move-detect-edge"
|
||||
version = "0.1.0"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = ["numpy", "opencv-python", "requests"]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = ["pytest"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
addopts = "-q"
|
||||
21
tests/test_models.py
Normal file
21
tests/test_models.py
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
from edge.models import Keypoint, BBox, Person, L_SHOULDER, R_HIP, CORE_JOINTS
|
||||
|
||||
def test_bbox_derived_properties():
|
||||
b = BBox(10, 20, 30, 80)
|
||||
assert b.width == 20
|
||||
assert b.height == 60
|
||||
assert b.aspect == 20 / 60
|
||||
assert b.cx == 20 and b.cy == 50
|
||||
|
||||
def test_bbox_zero_height_aspect_is_zero():
|
||||
assert BBox(0, 0, 10, 0).aspect == 0.0
|
||||
|
||||
def test_person_holds_17_keypoints():
|
||||
kps = [Keypoint(0, 0, 0.9) for _ in range(17)]
|
||||
p = Person(BBox(0, 0, 1, 1), kps)
|
||||
assert len(p.keypoints) == 17
|
||||
assert p.keypoints[L_SHOULDER].conf == 0.9
|
||||
|
||||
def test_coco_index_constants():
|
||||
assert CORE_JOINTS == (5, 6, 11, 12)
|
||||
assert R_HIP == 12
|
||||
Loading…
Reference in a new issue