220 lines
12 KiB
Python
220 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import time
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import cv2
|
|
|
|
from .core import normalize_bbox, normalize_keypoints
|
|
|
|
|
|
def _elapsed_ns(start: int) -> float:
|
|
return (time.perf_counter_ns() - start) / 1e6
|
|
|
|
|
|
def _cuda_timed(call):
|
|
torch.cuda.synchronize()
|
|
started = time.perf_counter_ns()
|
|
result = call()
|
|
torch.cuda.synchronize()
|
|
return result, _elapsed_ns(started)
|
|
|
|
|
|
def _ort_iobinding_run(session, input_array: np.ndarray):
|
|
"""Measure H2D, graph execution, and D2H as separate synchronized stages."""
|
|
import onnxruntime as ort
|
|
device_input, h2d_ms = _cuda_timed(lambda: ort.OrtValue.ortvalue_from_numpy(input_array, "cuda", 0))
|
|
binding = session.io_binding()
|
|
binding.bind_ortvalue_input(session.get_inputs()[0].name, device_input)
|
|
for output in session.get_outputs():
|
|
binding.bind_output(output.name, "cuda", 0)
|
|
_, forward_ms = _cuda_timed(lambda: session.run_with_iobinding(binding))
|
|
device_outputs = binding.get_outputs()
|
|
outputs, d2h_ms = _cuda_timed(lambda: [value.numpy() for value in device_outputs])
|
|
return outputs, h2d_ms, forward_ms, d2h_ms
|
|
|
|
|
|
class Models:
|
|
"""Persistent model holder for both baseline and RTMLib execution paths."""
|
|
|
|
def __init__(self) -> None:
|
|
self.backend = os.getenv("MODEL_BACKEND", "openmmlab")
|
|
self.det_config = os.getenv("DET_CONFIG", "")
|
|
self.det_checkpoint = os.getenv("DET_CHECKPOINT", "")
|
|
self.pose_config = os.getenv("POSE_CONFIG", "")
|
|
self.pose_checkpoint = os.getenv("POSE_CHECKPOINT", "")
|
|
self.detector = self.pose = None
|
|
self.providers: dict[str, list[str]] = {}
|
|
self.input_metadata: dict[str, Any] = {}
|
|
|
|
@staticmethod
|
|
def sha256(path: str) -> str:
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as fh:
|
|
for chunk in iter(lambda: fh.read(8 * 1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
def load(self) -> None:
|
|
cv2.setNumThreads(int(os.getenv("CV2_NUM_THREADS", "2")))
|
|
torch.backends.cudnn.benchmark = True
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
if self.backend == "openmmlab":
|
|
from mmdet.apis import init_detector
|
|
from mmpose.apis import init_model
|
|
self.detector = init_detector(self.det_config, self.det_checkpoint, device="cuda:0")
|
|
self.detector.test_cfg.nms.iou_threshold = float(os.getenv("NMS_IOU_THRESHOLD", ".60"))
|
|
self.pose = init_model(self.pose_config, self.pose_checkpoint, device="cuda:0")
|
|
self.providers = {"available": ["PyTorch CUDA"], "detector": ["PyTorch CUDA"], "pose": ["PyTorch CUDA"]}
|
|
self.input_metadata = {
|
|
"detector_parameter_device": str(next(self.detector.parameters()).device),
|
|
"pose_parameter_device": str(next(self.pose.parameters()).device),
|
|
}
|
|
return
|
|
|
|
if self.backend not in ("rtmlib_body", "rtmlib_wholebody", "rtmlib_rtmw_l"):
|
|
raise ValueError(f"unsupported MODEL_BACKEND={self.backend}")
|
|
import onnxruntime as ort
|
|
from rtmlib import RTMPose, YOLOX
|
|
self.detector = YOLOX(self.det_checkpoint, model_input_size=(640, 640), mode="human",
|
|
nms_thr=float(os.getenv("NMS_IOU_THRESHOLD", ".60")),
|
|
score_thr=float(os.getenv("DETECTION_THRESHOLD", ".35")),
|
|
backend="onnxruntime", device="cuda:0")
|
|
self.pose = RTMPose(self.pose_checkpoint, model_input_size=(192, 256),
|
|
backend="onnxruntime", device="cuda:0", to_openpose=False)
|
|
self.providers = {
|
|
"available": ort.get_available_providers(),
|
|
"detector": self.detector.session.get_providers(),
|
|
"pose": self.pose.session.get_providers(),
|
|
}
|
|
for name, providers in self.providers.items():
|
|
if name != "available" and (not providers or providers[0] != "CUDAExecutionProvider"):
|
|
raise RuntimeError(f"{name} is not using CUDAExecutionProvider first: {providers}")
|
|
self.input_metadata = {
|
|
"detector": [{"name": x.name, "shape": x.shape, "type": x.type} for x in self.detector.session.get_inputs()],
|
|
"pose": [{"name": x.name, "shape": x.shape, "type": x.type} for x in self.pose.session.get_inputs()],
|
|
}
|
|
print(f"ort.get_available_providers={self.providers['available']}", flush=True)
|
|
print(f"detector_session.get_providers={self.providers['detector']}", flush=True)
|
|
print(f"pose_session.get_providers={self.providers['pose']}", flush=True)
|
|
|
|
def infer(self, frame: np.ndarray, threshold: float, max_people: int) -> tuple[list[dict], dict[str, float]]:
|
|
if self.backend == "openmmlab":
|
|
return self._infer_openmmlab(frame, threshold, max_people)
|
|
return self._infer_rtmlib(frame, threshold, max_people)
|
|
|
|
@torch.inference_mode()
|
|
def _infer_openmmlab(self, frame: np.ndarray, threshold: float, max_people: int):
|
|
from mmdet.apis import inference_detector
|
|
from mmengine.registry import init_default_scope
|
|
from mmpose.apis import inference_topdown
|
|
height, width = frame.shape[:2]
|
|
timing = {k: 0.0 for k in ("frame_copy_ms", "color_conversion_ms", "resize_letterbox_ms",
|
|
"cpu_to_gpu_ms", "detector_postprocess_ms", "bbox_crop_affine_ms",
|
|
"pose_decode_ms", "gpu_to_cpu_ms")}
|
|
init_default_scope("mmdet")
|
|
result, timing["detector_forward_ms"] = _cuda_timed(lambda: inference_detector(self.detector, frame))
|
|
t = time.perf_counter_ns()
|
|
pred = result.pred_instances.cpu().numpy()
|
|
keep = (pred.labels == 0) & (pred.scores >= threshold)
|
|
boxes, scores = pred.bboxes[keep], pred.scores[keep]
|
|
if len(boxes):
|
|
order = np.argsort(scores)[::-1][:max_people]
|
|
boxes, scores = boxes[order], scores[order]
|
|
timing["detector_postprocess_ms"] = _elapsed_ns(t)
|
|
init_default_scope("mmpose")
|
|
if len(boxes):
|
|
pose_results, timing["pose_forward_ms"] = _cuda_timed(lambda: inference_topdown(self.pose, frame, bboxes=boxes))
|
|
else:
|
|
pose_results, timing["pose_forward_ms"] = [], 0.0
|
|
t = time.perf_counter_ns()
|
|
people = []
|
|
for idx, (box, det_score, sample) in enumerate(zip(boxes, scores, pose_results)):
|
|
points = sample.pred_instances.keypoints[0]
|
|
kp_scores = sample.pred_instances.keypoint_scores[0]
|
|
combined = np.column_stack((points, kp_scores))
|
|
wholebody = normalize_keypoints(combined, width, height, 133)
|
|
people.append(self._person(idx, box, det_score, wholebody, "coco_wholebody_133", width, height))
|
|
timing["pose_decode_ms"] = _elapsed_ns(t)
|
|
timing["detector_ms"] = timing["detector_forward_ms"] + timing["detector_postprocess_ms"]
|
|
timing["pose_ms"] = timing["pose_forward_ms"] + timing["pose_decode_ms"]
|
|
return people, timing
|
|
|
|
def _yolox_postprocess(self, outputs: np.ndarray, ratio: float, threshold: float):
|
|
from rtmlib.tools.object_detection.post_processings import multiclass_nms
|
|
if outputs.shape[-1] == 5:
|
|
boxes, scores = outputs[0, :, :4] / ratio, outputs[0, :, 4]
|
|
keep = scores >= threshold
|
|
boxes, scores = boxes[keep], scores[keep]
|
|
order = np.argsort(scores)[::-1]
|
|
return boxes[order], scores[order]
|
|
grids, expanded = [], []
|
|
for stride in (8, 16, 32):
|
|
h = self.detector.model_input_size[0] // stride
|
|
w = self.detector.model_input_size[1] // stride
|
|
xv, yv = np.meshgrid(np.arange(w), np.arange(h))
|
|
grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
|
|
grids.append(grid); expanded.append(np.full((*grid.shape[:2], 1), stride))
|
|
outputs = outputs.copy()
|
|
outputs[..., :2] = (outputs[..., :2] + np.concatenate(grids, 1)) * np.concatenate(expanded, 1)
|
|
outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * np.concatenate(expanded, 1)
|
|
pred = outputs[0]; boxes = pred[:, :4]; scores = pred[:, 4:5] * pred[:, 5:]
|
|
xyxy = np.empty_like(boxes)
|
|
xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2; xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
|
|
xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2; xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
|
|
dets, _ = multiclass_nms(xyxy / ratio, scores, nms_thr=float(os.getenv("NMS_IOU_THRESHOLD", ".60")), score_thr=threshold)
|
|
if dets is None:
|
|
return np.empty((0, 4)), np.empty((0,))
|
|
dets = dets[dets[:, 5].astype(int) == 0]
|
|
order = np.argsort(dets[:, 4])[::-1]
|
|
return dets[order, :4], dets[order, 4]
|
|
|
|
def _infer_rtmlib(self, frame: np.ndarray, threshold: float, max_people: int):
|
|
height, width = frame.shape[:2]
|
|
timing = {k: 0.0 for k in ("frame_copy_ms", "color_conversion_ms", "cpu_to_gpu_ms", "gpu_to_cpu_ms")}
|
|
t = time.perf_counter_ns(); det_image, ratio = self.detector.preprocess(frame)
|
|
timing["resize_letterbox_ms"] = _elapsed_ns(t)
|
|
det_input = np.ascontiguousarray(det_image.transpose(2, 0, 1), dtype=np.float32)[None]
|
|
self.input_metadata["detector_runtime"] = {"device": "CPU pinned by ORT", "dtype": str(det_input.dtype), "shape": list(det_input.shape)}
|
|
det_outputs, timing["cpu_to_gpu_ms"], timing["detector_forward_ms"], timing["gpu_to_cpu_ms"] = _ort_iobinding_run(self.detector.session, det_input)
|
|
det_output = det_outputs[0]
|
|
t = time.perf_counter_ns(); boxes, det_scores = self._yolox_postprocess(det_output, ratio, threshold)
|
|
boxes, det_scores = boxes[:max_people], det_scores[:max_people]
|
|
timing["detector_postprocess_ms"] = _elapsed_ns(t)
|
|
people = []
|
|
pose_forward = pose_decode = crop_affine = 0.0
|
|
for idx, (box, det_score) in enumerate(zip(boxes, det_scores)):
|
|
t = time.perf_counter_ns(); pose_image, center, scale = self.pose.preprocess(frame, box)
|
|
crop_affine += _elapsed_ns(t)
|
|
pose_input = np.ascontiguousarray(pose_image.transpose(2, 0, 1), dtype=np.float32)[None]
|
|
self.input_metadata["pose_runtime"] = {"device": "CPU pinned by ORT", "dtype": str(pose_input.dtype), "shape": list(pose_input.shape)}
|
|
outputs, h2d, elapsed, d2h = _ort_iobinding_run(self.pose.session, pose_input)
|
|
timing["cpu_to_gpu_ms"] += h2d; timing["gpu_to_cpu_ms"] += d2h; pose_forward += elapsed
|
|
t = time.perf_counter_ns(); points, scores = self.pose.postprocess(outputs, center, scale)
|
|
pose_decode += _elapsed_ns(t)
|
|
combined = np.column_stack((points[0], scores[0]))
|
|
count = 17 if self.backend == "rtmlib_body" else 133
|
|
normalized = normalize_keypoints(combined, width, height, count)
|
|
people.append(self._person(idx, box, det_score, normalized,
|
|
"coco17" if count == 17 else "coco_wholebody_133", width, height))
|
|
timing.update({"bbox_crop_affine_ms": crop_affine, "pose_forward_ms": pose_forward,
|
|
"pose_decode_ms": pose_decode,
|
|
"detector_ms": timing["detector_forward_ms"] + timing["detector_postprocess_ms"],
|
|
"pose_ms": crop_affine + pose_forward + pose_decode})
|
|
return people, timing
|
|
|
|
@staticmethod
|
|
def _person(idx, box, det_score, keypoints, schema, width, height):
|
|
body17 = keypoints[:17]
|
|
item = {"detection_id": idx, "bbox_xyxy_norm": normalize_bbox(box, width, height),
|
|
"det_score": float(np.clip(det_score, 0, 1)), "keypoint_schema": schema,
|
|
"body17": body17, "pose_score": float(np.mean([p[2] for p in body17]))}
|
|
if len(keypoints) == 133:
|
|
item["keypoints_norm"] = keypoints
|
|
return item
|