76 lines
3.1 KiB
Python
76 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import time
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from .core import normalize_bbox, normalize_keypoints
|
|
|
|
|
|
class Models:
|
|
def __init__(self) -> None:
|
|
self.det_config = os.environ["DET_CONFIG"]
|
|
self.det_checkpoint = os.environ["DET_CHECKPOINT"]
|
|
self.pose_config = os.environ["POSE_CONFIG"]
|
|
self.pose_checkpoint = os.environ["POSE_CHECKPOINT"]
|
|
self.detector = self.pose = None
|
|
|
|
@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:
|
|
from mmdet.apis import init_detector
|
|
from mmpose.apis import init_model
|
|
torch.backends.cudnn.benchmark = True
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
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")
|
|
|
|
@torch.inference_mode()
|
|
def infer(self, frame: np.ndarray, threshold: float, max_people: int) -> tuple[list[dict], float, float]:
|
|
from mmdet.apis import inference_detector
|
|
from mmengine.registry import init_default_scope
|
|
from mmpose.apis import inference_topdown
|
|
height, width = frame.shape[:2]
|
|
t0 = time.perf_counter_ns()
|
|
init_default_scope("mmdet")
|
|
result = inference_detector(self.detector, frame)
|
|
pred = result.pred_instances.cpu().numpy()
|
|
keep = (pred.labels == 0) & (pred.scores >= threshold)
|
|
boxes = pred.bboxes[keep]
|
|
scores = pred.scores[keep]
|
|
if len(boxes):
|
|
order = np.argsort(scores)[::-1][:max_people]
|
|
boxes, scores = boxes[order], scores[order]
|
|
det_ms = (time.perf_counter_ns() - t0) / 1e6
|
|
t1 = time.perf_counter_ns()
|
|
init_default_scope("mmpose")
|
|
pose_results = inference_topdown(self.pose, frame, bboxes=boxes) if len(boxes) else []
|
|
pose_ms = (time.perf_counter_ns() - t1) / 1e6
|
|
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({
|
|
"detection_id": idx,
|
|
"bbox_xyxy_norm": normalize_bbox(box, width, height),
|
|
"det_score": float(det_score),
|
|
"keypoint_schema": "coco_wholebody_133",
|
|
"keypoints_norm": wholebody,
|
|
"body17": wholebody[:17],
|
|
"pose_score": float(np.mean(np.nan_to_num(kp_scores[:17]))),
|
|
})
|
|
return people, det_ms, pose_ms
|