6c8fd9318f
Replace the NotImplementedError stub with a working HailoPoseAdapter for Pi5 + Hailo-8. yolov8s_pose has no on-chip NMS, so the adapter does host-side DFL box decode + keypoint decode + greedy NMS + letterbox un-map + COCO-17 assembly (numpy decode ported from Hailo's MIT PoseEstPostProcessing, locked by a synthetic 9-tensor unit test; the HailoRT device wiring is gated + device-validated on bringup). daemon.main() gains a POSE_BACKEND switch (hailo|movenet).
198 lines
9 KiB
Python
198 lines
9 KiB
Python
"""Hailo yolov8s_pose HEF adapter (실장비 Pi5 + Hailo-8/8L). Implements PosePort.
|
|
|
|
yolov8s_pose has NO on-chip NMS: the HEF emits 9 raw feature-map tensors (3 strides
|
|
x {box DFL (H,W,64), class score (H,W,1, already sigmoid), keypoints (H,W,51)}). This
|
|
adapter does the host-side DFL box-decode + keypoint-decode + NMS and remaps to
|
|
COCO-17 `list[Person]`.
|
|
|
|
The pure-numpy decode below is ported from Hailo's MIT-licensed reference
|
|
(Hailo-Application-Code-Examples/runtime/python/pose_estimation/pose_estimation_utils.py,
|
|
class PoseEstPostProcessing) and locked by tests/test_hailo_pose.py via an injected
|
|
fake infer_fn (no hardware needed).
|
|
|
|
DEVICE VALIDATION (bringup): the HailoRT wiring in `_build_hailo_infer` cannot run off
|
|
Pi/hailort. Validate it on device per docs/hardware-bringup.md, and confirm the model's
|
|
input (640x640 uint8 RGB, on-chip normalization) and the 9-tensor output shapes against
|
|
the actual HEF before trusting inference. The remap/decode math is what the unit tests fix.
|
|
"""
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from edge.models import Person, BBox, Keypoint
|
|
|
|
# COCO-17 keypoint order (yolov8_pose): nose, l/r eye, l/r ear, l/r shoulder,
|
|
# l/r elbow, l/r wrist, l/r hip, l/r knee, l/r ankle. Emitted in this order, so the
|
|
# 17 keypoints map positionally to edge.models.Keypoint without reindexing.
|
|
_NUM_KPTS = 17
|
|
|
|
|
|
def _sigmoid(x):
|
|
return 1.0 / (1.0 + np.exp(-x))
|
|
|
|
|
|
def _softmax(x):
|
|
e = np.exp(x - np.max(x, axis=-1, keepdims=True))
|
|
return e / np.sum(e, axis=-1, keepdims=True)
|
|
|
|
|
|
def _letterbox(frame, size):
|
|
"""Aspect-preserving resize to size x size with centered padding; BGR->RGB uint8.
|
|
Returns (canvas, scale, (pad_x, pad_y)) so results can be mapped back to `frame`."""
|
|
h, w = frame.shape[:2]
|
|
scale = min(size / w, size / h)
|
|
nw, nh = max(1, round(w * scale)), max(1, round(h * scale))
|
|
resized = cv2.resize(frame, (nw, nh))
|
|
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
|
|
canvas = np.full((size, size, 3), 114, dtype=np.uint8)
|
|
pad_x, pad_y = (size - nw) // 2, (size - nh) // 2
|
|
canvas[pad_y:pad_y + nh, pad_x:pad_x + nw] = rgb
|
|
return canvas, scale, (pad_x, pad_y)
|
|
|
|
|
|
def _pick(outs, size, channels):
|
|
for t in outs:
|
|
if t.shape[1] == size and t.shape[2] == size and t.shape[3] == channels:
|
|
return t
|
|
raise ValueError(f"missing yolov8s_pose tensor for size={size} channels={channels}")
|
|
|
|
|
|
def _iou(box, boxes):
|
|
x1 = np.maximum(box[0], boxes[:, 0]); y1 = np.maximum(box[1], boxes[:, 1])
|
|
x2 = np.minimum(box[2], boxes[:, 2]); y2 = np.minimum(box[3], boxes[:, 3])
|
|
inter = np.clip(x2 - x1, 0, None) * np.clip(y2 - y1, 0, None)
|
|
area = (box[2] - box[0]) * (box[3] - box[1])
|
|
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
|
|
union = area + areas - inter
|
|
return np.where(union > 0, inter / union, 0.0)
|
|
|
|
|
|
def _nms(pred, conf_thres, iou_thres, max_det):
|
|
"""pred: (N, 4 + 1 + 51) = [xywh, score, 17*(x,y,score)]. Greedy IoU NMS (1 class)."""
|
|
keep_mask = pred[:, 4] > conf_thres
|
|
pred = pred[keep_mask]
|
|
if pred.shape[0] == 0:
|
|
return []
|
|
xywh = pred[:, :4]
|
|
xyxy = np.empty_like(xywh)
|
|
xyxy[:, 0] = xywh[:, 0] - xywh[:, 2] / 2
|
|
xyxy[:, 1] = xywh[:, 1] - xywh[:, 3] / 2
|
|
xyxy[:, 2] = xywh[:, 0] + xywh[:, 2] / 2
|
|
xyxy[:, 3] = xywh[:, 1] + xywh[:, 3] / 2
|
|
scores = pred[:, 4]
|
|
kpts = pred[:, 5:].reshape(-1, _NUM_KPTS, 3)
|
|
order = np.argsort(-scores)
|
|
dets = []
|
|
while order.size > 0 and len(dets) < max_det:
|
|
i = order[0]
|
|
dets.append({"bbox": xyxy[i], "score": float(scores[i]), "kpts": kpts[i]})
|
|
if order.size == 1:
|
|
break
|
|
ious = _iou(xyxy[i], xyxy[order[1:]])
|
|
order = order[1:][ious <= iou_thres]
|
|
return dets
|
|
|
|
|
|
def decode_yolov8_pose(outs, input_size, strides, reg_max,
|
|
conf_thres, iou_thres, max_det):
|
|
"""Host decode of the 9 raw yolov8s_pose tensors -> NMS'd detections (letterbox px).
|
|
Ported from Hailo's MIT PoseEstPostProcessing (decoder + extract + NMS)."""
|
|
reg_range = np.arange(reg_max + 1)
|
|
boxes_all = scores_all = kpts_all = None
|
|
for stride in strides:
|
|
size = input_size // stride
|
|
box = _pick(outs, size, 4 * (reg_max + 1)).astype(np.float32)
|
|
score = _pick(outs, size, 1).astype(np.float32)
|
|
kpt = _pick(outs, size, _NUM_KPTS * 3).astype(np.float32).copy()
|
|
|
|
gx, gy = np.meshgrid(np.arange(size) + 0.5, np.arange(size) + 0.5)
|
|
ct_col, ct_row = gx.flatten() * stride, gy.flatten() * stride
|
|
center = np.stack((ct_col, ct_row, ct_col, ct_row), axis=1) # (HW, 4)
|
|
|
|
bd = box.reshape(1, size * size, 4, reg_max + 1)
|
|
bd = _softmax(bd) * reg_range.reshape(1, 1, 1, -1)
|
|
bd = np.sum(bd, axis=-1) * stride # (1, HW, 4) distances (px)
|
|
bd = np.concatenate([bd[:, :, :2] * (-1), bd[:, :, 2:]], axis=-1)
|
|
decode_box = center[None] + bd # xyxy
|
|
xmin, ymin = decode_box[:, :, 0], decode_box[:, :, 1]
|
|
xmax, ymax = decode_box[:, :, 2], decode_box[:, :, 3]
|
|
xywh = np.transpose(np.array([(xmin + xmax) / 2, (ymin + ymax) / 2,
|
|
xmax - xmin, ymax - ymin]), [1, 2, 0])
|
|
|
|
sc = score.reshape(1, size * size, 1)
|
|
kp = kpt.reshape(1, size * size, _NUM_KPTS, 3)
|
|
kp[..., :2] *= 2
|
|
kp[..., :2] = stride * (kp[..., :2] - 0.5) + center[:, :2].reshape(size * size, 1, 2)
|
|
kp = kp.reshape(1, size * size, _NUM_KPTS * 3)
|
|
|
|
boxes_all = xywh if boxes_all is None else np.concatenate([boxes_all, xywh], axis=1)
|
|
scores_all = sc if scores_all is None else np.concatenate([scores_all, sc], axis=1)
|
|
kpts_all = kp if kpts_all is None else np.concatenate([kpts_all, kp], axis=1)
|
|
|
|
pred = np.concatenate([boxes_all, scores_all, kpts_all], axis=2)[0]
|
|
return _nms(pred, conf_thres, iou_thres, max_det)
|
|
|
|
|
|
class HailoPoseAdapter:
|
|
def __init__(self, hef_path, input_size=640, infer_fn=None,
|
|
score_threshold=0.5, nms_iou=0.7, reg_max=15, strides=(8, 16, 32),
|
|
min_person_score=0.2, max_detections=100):
|
|
self.hef_path = hef_path
|
|
self.input_size = input_size
|
|
self.score_threshold = score_threshold
|
|
self.nms_iou = nms_iou
|
|
self.reg_max = reg_max
|
|
self.strides = tuple(strides)
|
|
self.min_person_score = min_person_score
|
|
self.max_detections = max_detections
|
|
self._infer_fn = infer_fn if infer_fn is not None else self._build_hailo_infer(hef_path)
|
|
|
|
def _build_hailo_infer(self, hef_path): # pragma: no cover (device-only: needs hailort + HEF)
|
|
"""Lazily build the HailoRT inference closure. Validate on device (bringup).
|
|
Returns infer_fn(frame_640_uint8_rgb) -> list of the 9 output tensors."""
|
|
from hailo_platform import HEF, VDevice, FormatType, HailoSchedulingAlgorithm
|
|
|
|
params = VDevice.create_params()
|
|
params.scheduling_algorithm = HailoSchedulingAlgorithm.ROUND_ROBIN
|
|
vdevice = VDevice(params)
|
|
hef = HEF(hef_path)
|
|
infer_model = vdevice.create_infer_model(hef_path)
|
|
infer_model.set_batch_size(1)
|
|
infer_model.input().set_format_type(FormatType.UINT8) # on-chip normalization
|
|
out_infos = hef.get_output_vstream_infos()
|
|
for info in out_infos:
|
|
infer_model.output(info.name).set_format_type(FormatType.FLOAT32) # dequantized
|
|
configured = infer_model.configure()
|
|
self._hailo_ctx = (vdevice, configured) # keep alive for the session
|
|
|
|
def infer_fn(frame):
|
|
buffers = {info.name: np.empty(infer_model.output(info.name).shape, dtype=np.float32)
|
|
for info in out_infos}
|
|
bindings = configured.create_bindings(output_buffers=buffers)
|
|
bindings.input().set_buffer(np.ascontiguousarray(frame))
|
|
configured.run([bindings], timeout_ms=10000)
|
|
return [buffers[info.name][None, ...] if buffers[info.name].ndim == 3
|
|
else buffers[info.name] for info in out_infos]
|
|
|
|
return infer_fn
|
|
|
|
def infer(self, frame) -> list:
|
|
h, w = frame.shape[:2]
|
|
canvas, scale, (pad_x, pad_y) = _letterbox(frame, self.input_size)
|
|
outs = self._infer_fn(canvas)
|
|
dets = decode_yolov8_pose(outs, self.input_size, self.strides, self.reg_max,
|
|
self.score_threshold, self.nms_iou, self.max_detections)
|
|
persons = []
|
|
for d in dets:
|
|
if d["score"] < self.min_person_score:
|
|
continue
|
|
x1, y1, x2, y2 = d["bbox"]
|
|
bbox = BBox((x1 - pad_x) / scale, (y1 - pad_y) / scale,
|
|
(x2 - pad_x) / scale, (y2 - pad_y) / scale)
|
|
keypoints = []
|
|
for j in range(_NUM_KPTS):
|
|
kx, ky, ksc = d["kpts"][j]
|
|
keypoints.append(Keypoint(float((kx - pad_x) / scale),
|
|
float((ky - pad_y) / scale),
|
|
float(_sigmoid(ksc))))
|
|
persons.append(Person(bbox, keypoints))
|
|
return persons
|