34 lines
1.4 KiB
Python
34 lines
1.4 KiB
Python
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)]
|