217 lines
6.6 KiB
Python
217 lines
6.6 KiB
Python
from typing import List, Dict
|
|
import cv2
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
|
|
|
|
def letterbox(
|
|
img: np.ndarray,
|
|
new_shape=(960, 960),
|
|
color=(114, 114, 114),
|
|
auto=False,
|
|
scale_fill=False,
|
|
scale_up=True,
|
|
stride=32,
|
|
):
|
|
shape = img.shape[:2] # h, w
|
|
|
|
if isinstance(new_shape, int):
|
|
new_shape = (new_shape, new_shape)
|
|
|
|
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
|
if not scale_up:
|
|
r = min(r, 1.0)
|
|
|
|
ratio = (r, r)
|
|
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]
|
|
|
|
if auto:
|
|
dw, dh = np.mod(dw, stride), np.mod(dh, stride)
|
|
elif scale_fill:
|
|
dw, dh = 0.0, 0.0
|
|
new_unpad = (new_shape[1], new_shape[0])
|
|
ratio = (new_shape[1] / shape[1], new_shape[0] / shape[0])
|
|
|
|
dw /= 2
|
|
dh /= 2
|
|
|
|
if shape[::-1] != new_unpad:
|
|
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
|
|
|
|
top = int(round(dh - 0.1))
|
|
bottom = int(round(dh + 0.1))
|
|
left = int(round(dw - 0.1))
|
|
right = int(round(dw + 0.1))
|
|
|
|
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
|
|
return img, ratio, (dw, dh)
|
|
|
|
|
|
def xywh2xyxy(x: np.ndarray) -> np.ndarray:
|
|
y = np.copy(x)
|
|
y[:, 0] = x[:, 0] - x[:, 2] / 2
|
|
y[:, 1] = x[:, 1] - x[:, 3] / 2
|
|
y[:, 2] = x[:, 0] + x[:, 2] / 2
|
|
y[:, 3] = x[:, 1] + x[:, 3] / 2
|
|
return y
|
|
|
|
|
|
def box_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
|
|
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.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
|
|
area1 = max(0, box[2] - box[0]) * max(0, box[3] - box[1])
|
|
area2 = np.maximum(0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0, boxes[:, 3] - boxes[:, 1])
|
|
union = area1 + area2 - inter + 1e-7
|
|
return inter / union
|
|
|
|
|
|
def nms_numpy(boxes: np.ndarray, scores: np.ndarray, iou_thres: float) -> List[int]:
|
|
if len(boxes) == 0:
|
|
return []
|
|
|
|
order = scores.argsort()[::-1]
|
|
keep = []
|
|
|
|
while order.size > 0:
|
|
i = order[0]
|
|
keep.append(int(i))
|
|
if order.size == 1:
|
|
break
|
|
ious = box_iou(boxes[i], boxes[order[1:]])
|
|
inds = np.where(ious <= iou_thres)[0]
|
|
order = order[inds + 1]
|
|
|
|
return keep
|
|
|
|
|
|
def scale_boxes_from_letterbox(boxes: np.ndarray, ratio, pad, original_shape) -> np.ndarray:
|
|
# boxes: xyxy in letterboxed image
|
|
gain = ratio[0]
|
|
pad_x, pad_y = pad
|
|
|
|
boxes[:, [0, 2]] -= pad_x
|
|
boxes[:, [1, 3]] -= pad_y
|
|
boxes[:, :4] /= gain
|
|
|
|
h, w = original_shape[:2]
|
|
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w - 1)
|
|
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h - 1)
|
|
return boxes
|
|
|
|
|
|
class YOLOv5ONNXDetector:
|
|
def __init__(
|
|
self,
|
|
model_path: str,
|
|
class_names: List[str],
|
|
img_size: int = 960,
|
|
device: str = "cpu",
|
|
):
|
|
self.model_path = model_path
|
|
self.class_names = class_names
|
|
self.img_size = img_size
|
|
|
|
providers = ["CPUExecutionProvider"]
|
|
if device.lower() in ["cuda", "gpu"]:
|
|
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
|
|
|
self.session = ort.InferenceSession(model_path, providers=providers)
|
|
active_providers = self.session.get_providers()
|
|
print("[YOLO] providers:", active_providers)
|
|
if (
|
|
device.lower() in ["cuda", "gpu"]
|
|
and "CUDAExecutionProvider" not in active_providers
|
|
):
|
|
raise RuntimeError(
|
|
"请求使用 CUDA 推理，但 YOLO 会话未启用 "
|
|
"CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 "
|
|
"onnxruntime-gpu。"
|
|
)
|
|
self.input_name = self.session.get_inputs()[0].name
|
|
self.output_names = [o.name for o in self.session.get_outputs()]
|
|
|
|
def preprocess(self, img_bgr: np.ndarray):
|
|
img_lb, ratio, pad = letterbox(img_bgr, new_shape=(self.img_size, self.img_size), auto=False)
|
|
img_rgb = cv2.cvtColor(img_lb, cv2.COLOR_BGR2RGB)
|
|
img = img_rgb.transpose(2, 0, 1).astype(np.float32) / 255.0
|
|
img = np.expand_dims(img, axis=0)
|
|
return img, ratio, pad
|
|
|
|
def infer(
|
|
self,
|
|
img_bgr: np.ndarray,
|
|
conf_thres: float = 0.25,
|
|
iou_thres: float = 0.45,
|
|
max_det: int = 50,
|
|
) -> List[Dict]:
|
|
input_tensor, ratio, pad = self.preprocess(img_bgr)
|
|
outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
|
|
|
|
pred = outputs[0]
|
|
|
|
# 常见 YOLOv5 ONNX 输出: [1, num_boxes, 5 + num_classes]
|
|
if isinstance(pred, list):
|
|
pred = pred[0]
|
|
pred = np.asarray(pred)
|
|
|
|
if pred.ndim == 3:
|
|
pred = pred[0]
|
|
elif pred.ndim == 2:
|
|
pass
|
|
else:
|
|
raise ValueError(f"不支持的 YOLO 输出维度: {pred.shape}")
|
|
|
|
if pred.shape[1] < 6:
|
|
raise ValueError(f"YOLO 输出最后一维应至少为 6，当前为: {pred.shape}")
|
|
|
|
boxes_xywh = pred[:, :4]
|
|
obj_conf = pred[:, 4]
|
|
|
|
if pred.shape[1] == 6:
|
|
# 某些导出格式可能已经是 x,y,w,h,conf,class_id
|
|
cls_ids = pred[:, 5].astype(np.int32)
|
|
scores = obj_conf
|
|
else:
|
|
cls_scores = pred[:, 5:]
|
|
cls_ids = np.argmax(cls_scores, axis=1)
|
|
cls_conf = cls_scores[np.arange(len(cls_scores)), cls_ids]
|
|
scores = obj_conf * cls_conf
|
|
|
|
mask = scores >= conf_thres
|
|
if not np.any(mask):
|
|
return []
|
|
|
|
boxes_xywh = boxes_xywh[mask]
|
|
scores = scores[mask]
|
|
cls_ids = cls_ids[mask]
|
|
|
|
boxes = xywh2xyxy(boxes_xywh)
|
|
boxes = scale_boxes_from_letterbox(boxes, ratio, pad, img_bgr.shape)
|
|
|
|
final_results = []
|
|
|
|
for class_id in np.unique(cls_ids):
|
|
idxs = np.where(cls_ids == class_id)[0]
|
|
keep = nms_numpy(boxes[idxs], scores[idxs], iou_thres)
|
|
|
|
for k in keep:
|
|
original_idx = idxs[k]
|
|
cid = int(class_id)
|
|
cname = self.class_names[cid] if 0 <= cid < len(self.class_names) else str(cid)
|
|
final_results.append({
|
|
"bbox_xyxy": boxes[original_idx].astype(float).tolist(),
|
|
"det_confidence": float(scores[original_idx]),
|
|
"det_class_id": cid,
|
|
"det_class_name": cname,
|
|
})
|
|
|
|
final_results.sort(key=lambda x: x["det_confidence"], reverse=True)
|
|
return final_results[:max_det]
|