feat(edge): MoveNet ONNX CPU pose adapter (COCO-17 remap)
This commit is contained in:
parent
581f337021
commit
d67513fc0e
4 changed files with 82 additions and 0 deletions
34
edge/adapters/movenet_pose.py
Normal file
34
edge/adapters/movenet_pose.py
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from edge.models import Person, BBox, Keypoint
|
||||
|
||||
|
||||
class MoveNetPoseAdapter:
|
||||
"""MoveNet SinglePose Lightning (COCO-17) via onnxruntime. CPU dev/pilot PosePort.
|
||||
출력 [1,1,17,3] = (y, x, score) normalized 0..1."""
|
||||
|
||||
def __init__(self, onnx_path: str, input_size: int = 192, session=None,
|
||||
min_person_score: float = 0.2):
|
||||
self.input_size = input_size
|
||||
self.min_person_score = min_person_score
|
||||
if session is None:
|
||||
import onnxruntime # gated: bundled on build machine
|
||||
session = onnxruntime.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
|
||||
self.session = session
|
||||
self._input_name = session.get_inputs()[0].name
|
||||
|
||||
def infer(self, frame) -> list:
|
||||
h, w = frame.shape[:2]
|
||||
img = cv2.resize(frame, (self.input_size, self.input_size))
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.int32)[None, ...]
|
||||
out = self.session.run(None, {self._input_name: img})[0] # [1,1,17,3]
|
||||
kp = out[0, 0] # [17,3] (y,x,score)
|
||||
keypoints = [Keypoint(float(kp[j, 1] * w), float(kp[j, 0] * h), float(kp[j, 2]))
|
||||
for j in range(17)]
|
||||
vis = [k for k in keypoints if k.conf >= self.min_person_score]
|
||||
if not vis:
|
||||
return []
|
||||
xs = [k.x for k in vis]; ys = [k.y for k in vis]
|
||||
bbox = BBox(min(xs), min(ys), max(xs), max(ys))
|
||||
return [Person(bbox, keypoints)]
|
||||
0
edge/tools/__init__.py
Normal file
0
edge/tools/__init__.py
Normal file
16
edge/tools/fetch_movenet.py
Normal file
16
edge/tools/fetch_movenet.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
"""빌드머신 전용: MoveNet SinglePose Lightning ONNX 다운로드(폐쇄망 번들). Apache-2.0."""
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
# MoveNet Lightning ONNX (예시 소스; 조직 미러로 교체 가능)
|
||||
URL = "https://github.com/PINTO0309/PINTO_model_zoo/raw/main/115_MoveNet/resources/movenet_singlepose_lightning_4.onnx"
|
||||
DEST = os.path.join(os.path.dirname(__file__), "..", "..", "models", "movenet_lightning.onnx")
|
||||
|
||||
def main():
|
||||
os.makedirs(os.path.dirname(DEST), exist_ok=True)
|
||||
print(f"downloading MoveNet ONNX -> {DEST}")
|
||||
urllib.request.urlretrieve(URL, DEST)
|
||||
print("done")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
32
tests/test_movenet_pose.py
Normal file
32
tests/test_movenet_pose.py
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
import numpy as np
|
||||
from edge.adapters.movenet_pose import MoveNetPoseAdapter
|
||||
|
||||
class _FakeSession:
|
||||
"""MoveNet 출력 형태 [1,1,17,3] (y,x,score) 정규화."""
|
||||
def get_inputs(self):
|
||||
class I: name = "input"
|
||||
return [I()]
|
||||
def run(self, _out, feed):
|
||||
kp = np.zeros((1, 1, 17, 3), dtype=np.float32)
|
||||
for j in range(17):
|
||||
kp[0, 0, j] = [0.5, 0.25, 0.9] # y=0.5, x=0.25, score=0.9
|
||||
return [kp]
|
||||
|
||||
def test_infer_maps_movenet_to_coco17_persons():
|
||||
frame = np.zeros((400, 800, 3), dtype=np.uint8)
|
||||
a = MoveNetPoseAdapter("unused.onnx", session=_FakeSession())
|
||||
persons = a.infer(frame)
|
||||
assert len(persons) == 1
|
||||
p = persons[0]
|
||||
assert len(p.keypoints) == 17
|
||||
# x = 0.25 * 800 = 200, y = 0.5 * 400 = 200
|
||||
assert abs(p.keypoints[0].x - 200) < 1 and abs(p.keypoints[0].y - 200) < 1
|
||||
assert abs(p.keypoints[0].conf - 0.9) < 1e-6
|
||||
assert p.bbox.x1 <= 200 <= p.bbox.x2
|
||||
|
||||
def test_low_score_person_dropped():
|
||||
class Low(_FakeSession):
|
||||
def run(self, _o, feed):
|
||||
return [np.zeros((1, 1, 17, 3), dtype=np.float32)] # all score 0
|
||||
a = MoveNetPoseAdapter("unused.onnx", session=Low())
|
||||
assert a.infer(np.zeros((10, 10, 3), np.uint8)) == []
|
||||
Loading…
Reference in a new issue