278 lines
9.3 KiB
Python
278 lines
9.3 KiB
Python
# npu_yolo_pose_onnx.py
|
||
# 修复要点：
|
||
# 1. 正确处理 YOLOv8 Pose anchor 输出（避免 40+ 人）
|
||
# 2. 关键点坐标正确逆 letterbox（减 padding 再除 ratio）
|
||
# 3. visibility 使用 sigmoid
|
||
# 4. NMS 后限制最大人数，保证工程稳定性
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import onnxruntime as ort
|
||
|
||
from utils.logger import get_logger
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
# -------------------------------------------------
|
||
# Letterbox
|
||
# -------------------------------------------------
|
||
def letterbox(img, new_shape=(1280, 1280), color=(114, 114, 114)):
|
||
shape = img.shape[:2] # h, w
|
||
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||
|
||
new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
|
||
dw = new_shape[1] - new_unpad[0]
|
||
dh = new_shape[0] - new_unpad[1]
|
||
dw /= 2
|
||
dh /= 2
|
||
|
||
if shape[::-1] != new_unpad:
|
||
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||
|
||
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
|
||
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
|
||
img = cv2.copyMakeBorder(
|
||
img, top, bottom, left, right,
|
||
cv2.BORDER_CONSTANT, value=color
|
||
)
|
||
return img, r, (dw, dh)
|
||
|
||
|
||
# -------------------------------------------------
|
||
# Pose Skeleton Definition (COCO-17)
|
||
# -------------------------------------------------
|
||
POSE_SKELETON = [
|
||
(16,14),(14,12),(17,15),(15,13),(12,13),
|
||
(6,12),(7,13),(6,7),
|
||
(6,8),(7,9),(8,10),(9,11),
|
||
(2,3),(1,2),(1,3),
|
||
(2,4),(3,5),(4,6),(5,7)
|
||
]
|
||
|
||
POSE_SKELETON = [(a-1, b-1) for (a, b) in POSE_SKELETON]
|
||
|
||
POSE_COLORS = [
|
||
(255,0,0),(255,85,0),(255,170,0),(255,255,0),
|
||
(170,255,0),(85,255,0),(0,255,0),
|
||
(0,255,85),(0,255,170),(0,255,255),
|
||
(0,170,255),(0,85,255),(0,0,255),
|
||
(85,0,255),(170,0,255),(255,0,255),(255,0,170)
|
||
]
|
||
|
||
|
||
# -------------------------------------------------
|
||
# YOLOv8 Pose ONNX
|
||
# -------------------------------------------------
|
||
class YOLOv8_Pose_ONNX:
|
||
def __init__(
|
||
self,
|
||
onnx_path,
|
||
conf_threshold=0.6, # ★ 提高阈值，避免 anchor 噪声
|
||
iou_threshold=0.45,
|
||
input_size=1280,
|
||
max_persons=5 # ★ 限制最大人数
|
||
):
|
||
providers = [
|
||
("CANNExecutionProvider", {
|
||
"device_id": 0,
|
||
"arena_extend_strategy": "kNextPowerOfTwo",
|
||
"npu_mem_limit": 16 * 1024 * 1024 * 1024,
|
||
"precision_mode": "allow_fp32_to_fp16",
|
||
"op_select_impl_mode": "high_precision",
|
||
"enable_cann_graph": True,
|
||
}),
|
||
"CUDAExecutionProvider",
|
||
"CPUExecutionProvider",
|
||
]
|
||
|
||
self.session = ort.InferenceSession(onnx_path, providers=providers)
|
||
|
||
# 获取真实工作 provider
|
||
actual_providers = self.session.get_providers()
|
||
|
||
logger.info("YOLO Providers:", actual_providers)
|
||
|
||
if "CANNExecutionProvider" in actual_providers:
|
||
logger.info("[INFO] YOLO 使用 CANNExecutionProvider（昇腾）")
|
||
elif 'CUDAExecutionProvider' in actual_providers:
|
||
logger.info("[INFO] YOLO 使用 CUDAExecutionProvider（NVIDIA GPU）")
|
||
else:
|
||
logger.info("[INFO] YOLO 使用 CPUExecutionProvider（非昇腾环境）")
|
||
|
||
self.conf_threshold = conf_threshold
|
||
self.iou_threshold = iou_threshold
|
||
self.max_persons = max_persons
|
||
|
||
self.input_name = self.session.get_inputs()[0].name
|
||
self.input_size = (input_size, input_size)
|
||
logger.info(f"模型输入名称: {self.input_name}")
|
||
logger.info(f"模型输入形状: {self.session.get_inputs()[0].shape}")
|
||
logger.info(f"模型输出形状: {self.session.get_outputs()[0].shape}")
|
||
|
||
|
||
def nms(self, boxes, scores, iou_threshold=0.45):
|
||
"""
|
||
boxes: [N,4] xyxy
|
||
scores: [N]
|
||
"""
|
||
x1 = boxes[:, 0]
|
||
y1 = boxes[:, 1]
|
||
x2 = boxes[:, 2]
|
||
y2 = boxes[:, 3]
|
||
|
||
areas = (x2 - x1) * (y2 - y1)
|
||
order = scores.argsort()[::-1]
|
||
|
||
keep = []
|
||
while order.size > 0:
|
||
i = order[0]
|
||
keep.append(i)
|
||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||
|
||
w = np.maximum(0.0, xx2 - xx1)
|
||
h = np.maximum(0.0, yy2 - yy1)
|
||
inter = w * h
|
||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||
|
||
inds = np.where(ovr <= iou_threshold)[0]
|
||
order = order[inds + 1]
|
||
|
||
# 限制最大人数
|
||
if len(keep) >= self.max_persons:
|
||
break
|
||
|
||
return np.array(keep, dtype=int)
|
||
# -------------------------------------------------
|
||
def preprocess(self, img):
|
||
self.orig_shape = img.shape[:2]
|
||
img, self.ratio, (self.dw, self.dh) = letterbox(img, self.input_size)
|
||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||
img = img.transpose(2, 0, 1).astype(np.float32) / 255.0
|
||
img = np.expand_dims(img, axis=0)
|
||
return img
|
||
|
||
def postprocess(self, preds, im0_shape):
|
||
"""
|
||
preds: onnx output, shape = [1, 56, 33600]
|
||
im0_shape: (h, w) of original frame
|
||
"""
|
||
|
||
preds = preds[0] # [56, 33600]
|
||
preds = preds.transpose(1, 0) # [33600, 56]
|
||
|
||
# =============================
|
||
# 1. 拆分输出
|
||
# =============================
|
||
boxes = preds[:, 0:4] # cx, cy, w, h (input scale)
|
||
scores = preds[:, 4] # obj conf
|
||
kpts_raw = preds[:, 5:] # [33600, 51] = 17*3
|
||
|
||
# =============================
|
||
# 2. 置信度筛选
|
||
# =============================
|
||
mask = scores > self.conf_threshold
|
||
boxes = boxes[mask]
|
||
scores = scores[mask]
|
||
kpts_raw = kpts_raw[mask]
|
||
|
||
if boxes.shape[0] == 0:
|
||
return []
|
||
|
||
# =============================
|
||
# 3. bbox cxcywh -> xyxy（input scale）
|
||
# =============================
|
||
boxes_xyxy = np.zeros_like(boxes)
|
||
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2 # x1
|
||
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2 # y1
|
||
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2 # x2
|
||
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2 # y2
|
||
|
||
# =============================
|
||
# 4. inverse letterbox（bbox）
|
||
# =============================
|
||
boxes_xyxy[:, [0, 2]] = (boxes_xyxy[:, [0, 2]] - self.dw) / self.ratio
|
||
boxes_xyxy[:, [1, 3]] = (boxes_xyxy[:, [1, 3]] - self.dh) / self.ratio
|
||
|
||
boxes_xyxy[:, 0] = np.clip(boxes_xyxy[:, 0], 0, im0_shape[1])
|
||
boxes_xyxy[:, 1] = np.clip(boxes_xyxy[:, 1], 0, im0_shape[0])
|
||
boxes_xyxy[:, 2] = np.clip(boxes_xyxy[:, 2], 0, im0_shape[1])
|
||
boxes_xyxy[:, 3] = np.clip(boxes_xyxy[:, 3], 0, im0_shape[0])
|
||
|
||
# =============================
|
||
# 5. NMS
|
||
# =============================
|
||
keep = self.nms(boxes_xyxy, scores, self.iou_threshold)
|
||
boxes_xyxy = boxes_xyxy[keep]
|
||
scores = scores[keep]
|
||
kpts_raw = kpts_raw[keep]
|
||
|
||
# =============================
|
||
# 6. 逐人处理 keypoints（关键）
|
||
# =============================
|
||
results = []
|
||
|
||
for i in range(len(boxes_xyxy)):
|
||
x1, y1, x2, y2 = boxes_xyxy[i]
|
||
|
||
# (51,) -> (17,3)
|
||
kpts = kpts_raw[i].reshape(17, 3).copy()
|
||
|
||
|
||
kpts[:, 0] = (kpts[:, 0] - self.dw) / self.ratio
|
||
kpts[:, 1] = (kpts[:, 1] - self.dh) / self.ratio
|
||
|
||
# ✅ 加回 bbox offset（核心修复点）
|
||
#kpts[:, 0] += x1
|
||
#kpts[:, 1] += y1
|
||
|
||
# clip
|
||
kpts[:, 0] = np.clip(kpts[:, 0], 0, im0_shape[1])
|
||
kpts[:, 1] = np.clip(kpts[:, 1], 0, im0_shape[0])
|
||
|
||
# visibility sigmoid（防溢出）
|
||
kpts[:, 2] = 1.0 / (1.0 + np.exp(-np.clip(kpts[:, 2], -50, 50)))
|
||
|
||
results.append({
|
||
"bbox": [float(x1), float(y1), float(x2), float(y2)],
|
||
"conf": float(scores[i]),
|
||
"kpts": kpts
|
||
})
|
||
|
||
return results
|
||
|
||
# -------------------------------------------------
|
||
def __call__(self, frame):
|
||
inp = self.preprocess(frame)
|
||
# 耗时操作
|
||
pred = self.session.run(None, {self.input_name: inp})[0]
|
||
return self.postprocess(pred, frame.shape[:2])
|
||
|
||
@staticmethod
|
||
def draw_keypoints(frame, pose_results, vis_thres=0.3):
|
||
for res in pose_results:
|
||
kpts = res.get("kpts", None) # 注意这里对应 postprocess 返回的 key
|
||
if kpts is None or len(kpts) != 17:
|
||
continue
|
||
# 如果是 ndarray，转换为 list
|
||
if isinstance(kpts, np.ndarray):
|
||
kpts = kpts.tolist()
|
||
|
||
for i, (x, y, v) in enumerate(kpts):
|
||
if v > vis_thres:
|
||
cv2.circle(frame, (int(x), int(y)), 5, POSE_COLORS[i], -1)
|
||
|
||
for a, b in POSE_SKELETON:
|
||
if kpts[a][2] > vis_thres and kpts[b][2] > vis_thres:
|
||
cv2.line(
|
||
frame,
|
||
(int(kpts[a][0]), int(kpts[a][1])),
|
||
(int(kpts[b][0]), int(kpts[b][1])),
|
||
POSE_COLORS[a],
|
||
2
|
||
)
|
||
return frame
|
||
|