ncnn可用户移动端的识别，参考资料
Android工程中使用ncnn进行yolo识别-CSDN博客
系统: Windows11,显卡: RTX4060
一、训练集数据准备
目录结构如下
data.yaml
path: ./dataset # 数据集根目录的绝对路径或相对路径
train: images/train # 训练集图片相对路径
val: images/val # 验证集图片相对路径
nc: 12 # 类别数量
names: # 类别名称列表 (顺序必须与标注时的 class_id 对应!)
0: target1 # 对应你原来 train/0 的类别
1: target2
2: target3
3: target4
4: target5
5: target6
6: target7
7: target8
8: target9
9: target10
10: target11
11: target12图片标注信息
二、安装训练环境
conda创建激活环境
# 1. 创建环境（-n 指定名称，python=3.10 指定版本，-y 自动确认）
conda create -n ncnnenv python=3.10 -y
# 2. 激活环境
conda activate ncnnenvconda 安装 GPU 版 PyTorch（CUDA12.1）
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia安装完成后验证
python -c "import torch; print('cuda available:', torch.cuda.is_available()); print('GPU name:', torch.cuda.get_device_name(0))"预期输出：
cuda available: True
GPU name: NVIDIA GeForce RTX 4060 Laptop GPU三、安装训练环境
进入yolo_train.py所在目录
cd C:\Users\asus\Downloads\ncnn_train\yolo_train核对训练参数，用于移动端imgsz设为320
开始训练
python yolo_train.py 训练结束，生成best.pt和best.onnx
进入best.onnx所在目录
cd C:\Users\asus\Downloads\ncnn_train\yolo_train\runs\detect\runs\detect\gtx960_target_det\weights\导出ncnn
nnx "best.onnx" inputshape="[1,3,320,320]" device=cpu outname="my_yolov8_model" fp16=1四、在Android工程中使用
修改要加载的模块
修改输入尺寸为320
如果卡顿，可以隔几帧识别一次
五、参考原文及源代码
# YOLOv8 自定义目标检测模型 ------ 训练与部署操作说明
> 对应脚本：`yolo_train.py`
> 适用环境：Windows + NVIDIA GPU（参数针对 GTX 960 2G/4G 显存做过保守优化）
---
## 一、流程总览
脚本一次运行完成 4 个阶段：
```
环境检查 → 加载预训练模型 → 训练 → 验证评估 → 导出 ONNX（供 ncnn 部署）
```
| 阶段 | 做什么 | 产出 |
| ------- | ----------------------------- | --------------------------------------------- |
| 1. 环境检查 | 打印 PyTorch 版本、检测 CUDA/GPU | 无 GPU 直接报错退出 |
| 2. 加载模型 | 加载 `yolov8n.pt`（首次运行自动下载） | --- |
| 3. 训练 | 按 `dataset/data.yaml` 训练 10 轮 | `runs/detect/gtx960_target_det/`（含 `best.pt`） |
| 4. 验证 | 在验证集评估最佳模型 | 打印 mAP@50、mAP@50-95 |
| 5. 导出 | 导出 opset 12、固定尺寸 320 的 ONNX | `best.onnx`（供后续转 ncnn） |
---
## 二、目录结构约定
```
yolo_train/
├── yolo_train.py ← 训练主脚本
├── yolov8n.pt ← 预训练权重（首次自动下载）
├── dataset/
│ ├── data.yaml ← 数据集配置
│ ├── images/
│ │ ├── train/ ← 训练图片
│ │ └── val/ ← 验证图片
│ └── labels/
│ ├── train/ ← 训练标注（YOLO txt 格式）
│ └── val/
├── runs/detect/gtx960_target_det/ ← 训练结果输出目录
│ ├── weights/best.pt ← 最佳权重
│ ├── weights/last.pt ← 最后一轮权重
│ └── weights/best.onnx ← ONNX 模型
└── best.onnx ← 导出的 ONNX 模型
```
---
## 三、数据集配置（dataset/data.yaml）
- 当前测试共 **12 个类别**(按实际需要进行修改)：`target1` ~ `target12`
- **关键规则**：`names` 中的类别顺序（0~11）必须与标注时 `class_id` 一一对应，改错会导致检测类别错乱
- 图片路径为相对 `path` 的路径；首次训练会自动生成 `.cache` 文件，修改标注后若不生效请删除 `.cache`
---
## 四、运行方式
```bash
# 激活装有 ultralytics + GPU 版 torch 的环境后
python yolo_train.py
```
**依赖**：
```bash
pip install ultralytics
# torch 需安装 CUDA 版本（CPU 版会因检测不到 CUDA 而直接报错退出）
```
---
## 五、关键训练参数说明（GTX 960 显存优化）
| 参数 | 值 | 说明 |
| ---------- | ----- | --------------------------------------------------------------------------------------------------------- |
| `epochs` | 10 | 先跑 10 轮验证流程通不通，**确认无误后建议改为 100** |
| `imgsz` | 320 | 输入分辨率，GTX 960 不建议超过 640 |
| `batch` | 8 | 安全批次；**若爆显存（CUDA OOM）降为 4** |
| `workers` | 2 | 降低数据加载线程数，避免老平台内存溢出 |
| `cache` | False | 禁用图片缓存，防止显存被撑爆 |
| `amp` | False | 混合精度节省显存并提速（AMP 依赖 GPU 的 FP16 算力，GTX 960 是 Maxwell 架构（算力 5.2）当前代码值为 False，如需省显存提速可改 True，限RTX 20/30/40 系） |
| `device` | 0 | 使用 0 号 GPU |
| `exist_ok` | True | 允许覆盖同名实验目录，重复运行不会报错 |
**注意**：`if __name__ == "__main__"` 里的 `multiprocessing.freeze_support()` 是 Windows 下多进程数据加载的必需保护，请勿删除，也不要把训练逻辑挪到模块顶层。
---
## 六、训练结果解读
训练完成后查看 `runs/detect/gtx960_target_det/`：
- `weights/best.pt`：验证集表现最好的权重（一般用它做推理/导出）
- `weights/last.pt`：最后一轮权重
- `results.csv` / `results.png`：每轮 loss 与 mAP 曲线
- `confusion_matrix.png`：混淆矩阵，可查看哪些类别之间容易混淆
- 脚本最后会打印：
- `mAP@50`：IoU 阈值 0.5 下的平均精度
- `mAP@50-95`：更严格的综合指标
---
## 七、ONNX 导出与 ncnn 部署（后续步骤）
脚本导出的 ONNX 参数已为 ncnn 做适配：
- `opset=12`：ncnn 推荐版本
- `simplify=True`：简化计算图，提升推理速度
- `dynamic=False`：固定输入尺寸，ncnn 兼容性更好
**转 ncnn**（命令模板）：
```bash
pnnx "best.onnx" inputshape="[1,3,320,320]" device=cpu outname="my_yolov8_model" fp16=1
```
产出 `.param` + `.bin` 文件后，用 Python ncnn 加载：
```python
import ncnn
net = ncnn.Net()
net.load_param("best.ncnn.param")
net.load_model("best.ncnn.bin")
```
项目中已有的相关脚本可用于验证推理：
- `verify_onnx.py`：ONNX 推理验证
- `verify_ncnn.py`：ncnn 推理验证
其他改动位置：
```cpp
app/src/main/cpp/CustomPlayer.cpp
app/src/main/cpp/yolo.h
app/src/main/cpp/yolo.cpp
app/src/main/cpp/native-lib.cpp
```
---
## 八、常见问题（FAQ）
| 问题 | 原因/解决 |
| ------------------------ | ---------------------------------------------- |
| 启动即报"未检测到 CUDA 设备" | torch 装的是 CPU 版，需重装 CUDA 版 |
| `CUDA out of memory` | `batch` 降为 4，或 `workers` 降为 0/1 |
| 训练卡在数据加载 / Windows 多进程报错 | 确认保留 `freeze_support()`；可把 `workers` 设为 0 |
| 改了标注但训练没变化 | 删除 `dataset/labels/*.cache` 后重跑 |
| 类别检测错乱 | 检查 `data.yaml` 中 `names` 顺序与标注 `class_id` 是否一致 |
| 想正式训练 | 把 `epochs=10` 改为 100（或更高）后重跑 |
---
## 九、推荐操作流程
1. 准备图片与标注 → 放入 `dataset/images` 与 `dataset/labels`
2. 核对 `data.yaml` 的类别数与顺序
3. `python yolo_train.py` 跑 10 轮，确认流程无报错、loss 正常下降
4. 将 `epochs` 改为 100 正式训练
5. 查看 mAP 与混淆矩阵，必要时补充数据/调参
6. 取 `best.pt` 导出 ONNX → pnnx 转 ncnn → 用脚本验证推理yolo_train.py
python
import os
import torch
from ultralytics import YOLO
def main():
# ================= 环境检查 =================
print(f"🚀 PyTorch Version: {torch.__version__}")
print(f"🎮 CUDA Available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"💻 GPU Device: {torch.cuda.get_device_name(0)}")
print(f"📊 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024 ** 3:.1f} GB")
else:
raise RuntimeError("❌ 未检测到 CUDA 设备，请检查 PyTorch GPU 版本安装！")
# ================= 加载模型 =================
# 首次运行会自动下载 yolov8n.pt 预训练权重
model = YOLO("yolov8n.pt")
# ================= 启动训练 =================
# ⚠️ 以下参数已针对 GTX 960 (2G/4G显存) 进行保守优化
results = model.train(
data="dataset/data.yaml", # 数据集配置文件路径
epochs=100, # 💡 建议先跑10轮验证流程，确认无误后再改为100
imgsz=320, # 输入分辨率，GTX960 不建议超过640
batch=8, # 💡 GTX960 安全批次大小，爆显存则降为4
device=0, # 指定GPU编号
project="runs/detect", # 结果保存根目录
name="gtx960_target_det", # 本次实验名称
exist_ok=True, # 允许覆盖同名实验目录
workers=2, # 💡 降低线程数，避免老平台内存溢出
amp=False, # 开启混合精度，节省显存并提速
cache=False, # 💡 禁用缓存，防止GTX960显存被撑爆
verbose=True # 打印详细训练日志
)
# ================= 验证模型 =================
print("\n🔍 开始在验证集上评估最佳模型...")
metrics = model.val()
print(f"✅ mAP@50: {metrics.box.map50:.4f}")
print(f"✅ mAP@50-95: {metrics.box.map:.4f}")
# ================= 导出 ONNX =================
# 为后续转换 ncnn 做准备
print("\n📦 正在导出 ONNX 模型 (用于 ncnn 部署)...")
export_path = model.export(
format="onnx",
opset=12, # ncnn 推荐 opset 12
simplify=True, # 简化计算图，提升推理速度
dynamic=False, # 固定输入尺寸，ncnn 兼容性更好
imgsz=320
)
print(f"🎉 ONNX 模型已导出至: {export_path}")
if __name__ == "__main__":
# Windows 下多进程数据加载必须加此保护
import multiprocessing
multiprocessing.freeze_support()
main()
# 转换ncnn
# pnnx "F:\python.study\yolo_train\target_ring_sim.onnx" inputshape="[1,3,640,640]" device=cpu outname="my_yolov8_model" fp16=1
# import ncnn
#
# net = ncnn.Net()
# ret_param = net.load_param("target_ring_sim.ncnn.param")
# ret_model = net.load_model("target_ring_sim.ncnn.bin")
#
# print("load_param:", ret_param)
# print("load_model:", ret_model)verify_ncnn.py
python
import cv2
import ncnn
import numpy as np
import time
import os
# ================= 配置 =================
RTSP_URL = "rtsp://127.0.0.1:8554/stream/video"
PARAM_PATH = "best.ncnn.param"
BIN_PATH = "best.ncnn.bin"
INPUT_SIZE = 640
CONF_THRESH = 0.25
NMS_THRESH = 0.45
NAMES = [
"target1", "target2", "target3", "target4",
"target5", "target6", "target7", "target8",
"target9", "target10", "target11", "target12"
]
# ================= 兼容不同ncnn版本的API =================
def get_net_io_names(net):
"""尝试多种方式获取输入输出节点名，兼容不同ncnn python绑定版本"""
# 方式1: input_count() / get_input_layer() (旧版)
try:
n_in = net.input_count()
input_names = [net.get_input_layer(i).name for i in range(n_in)]
n_out = net.output_count()
output_names = [net.get_output_layer(i).name for i in range(n_out)]
print(f"[INFO] 使用 input_count()/get_input_layer() API")
return input_names, output_names
except AttributeError:
pass
# 方式2: inputs / outputs 属性 (部分新版)
try:
input_names = list(net.inputs)
output_names = list(net.outputs)
print(f"[INFO] 使用 net.inputs/outputs 属性")
return input_names, output_names
except (AttributeError, TypeError):
pass
# 方式3: 直接读param文件提取layer名
try:
input_names = []
output_names = []
with open(PARAM_PATH, 'r') as f:
lines = f.readlines()
for line in lines:
line = line.strip()
parts = line.split()
if len(parts) >= 2:
if parts[0] == 'Input':
input_names.append(parts[1])
elif parts[0] == 'Output':
output_names.append(parts[1])
if input_names or output_names:
print(f"[INFO] 从param文件解析到 layer名")
print(f" 输入: {input_names}")
print(f" 输出: {output_names}")
return input_names, output_names
except Exception as e:
print(f"[WARN] 读param文件失败: {e}")
# 方式4: 硬编码常见名称（YOLOv8 via pnnx 通常是 images / output0）
print(f"[WARN] 无法自动检测IO节点，使用pnnx默认名称: in0 -> out0")
return ["in0"], ["out0"]
# ================= 加载模型 =================
print("正在加载 ncnn 模型...")
net = ncnn.Net()
net.opt.use_vulkan_compute = True
net.opt.num_threads = 4
ret_p = net.load_param(PARAM_PATH)
ret_m = net.load_model(BIN_PATH)
print(f"load_param: {ret_p} load_model: {ret_m}")
input_names, output_names = get_net_io_names(net)
print(f"最终使用 -> 输入: {input_names}, 输出: {output_names}")
# ================= 打开RTSP流 =================
cap = cv2.VideoCapture(RTSP_URL, cv2.CAP_FFMPEG)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
if not cap.isOpened():
raise RuntimeError("RTSP 流打开失败")
# ================= 预处理（和ONNX版本一致的letterbox） =================
def preprocess_letterbox(image, size=640):
"""
和ONNX推理完全一致的letterbox预处理：
1. 等比缩放，保持宽高比
2. 灰边(padding=114)填充到640x640
3. BGR->RGB
4. 归一化到[0,1]
"""
h, w = image.shape[:2]
scale = size / max(h, w)
new_w = int(w * scale)
new_h = int(h * scale)
# 等比缩放
resized = cv2.resize(image, (new_w, new_h))
# 灰边填充（和ONNX中copyMakeBorder的value=(114,114,114)一致）
pad_w = (size - new_w) // 2
pad_h = (size - new_h) // 2
padded = cv2.copyMakeBorder(
resized,
pad_h, size - new_h - pad_h,
pad_w, size - new_w - pad_w,
cv2.BORDER_CONSTANT,
value=(114, 114, 114)
)
# BGR -> RGB
rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB)
# 用ncnn.Mat.from_pixels创建Mat
# 兼容不同版本：有些版本参数顺序不同
try:
mat = ncnn.Mat.from_pixels(rgb, ncnn.Mat.PixelType.PIXEL_RGB, size, size)
except TypeError:
# 某些版本可能是 from_pixels(data, w, h, type)
mat = ncnn.Mat.from_pixels(rgb, size, size, ncnn.Mat.PixelType.PIXEL_RGB)
# 归一化 [0,1]（和ONNX中 /255.0 一致）
mean = [0.0, 0.0, 0.0]
norm = [1 / 255.0, 1 / 255.0, 1 / 255.0]
try:
mat.substract_mean_normalize(mean, norm)
except AttributeError:
mat.substract_mean_normalize(mean, norm)
return mat, scale, pad_w, pad_h
def postprocess_letterbox_bug(outputs, scale, pad_w, pad_h, conf_thresh=0.25, iou_thresh=0.45):
"""
适配 pnnx 输出的 (16, 8400)：
前 4 维: cx, cy, w, h（640 尺度）
后 12 维: 12 类概率（已 sigmoid，不要再套一次）
"""
# outputs 是 ncnn.Mat，转numpy
data = np.array(outputs)
print(f"[DEBUG] ncnn output shape: {data.shape}, dtype: {data.dtype}")
# ============ 修改点 1：shape 分支重写 ============
if data.ndim == 3 and data.shape[0] == 1:
data = data[0]
if data.ndim != 2:
print(f"[ERROR] 无法识别的输出维度: {data.shape}")
return [], [], []
C, N = data.shape
num_classes = 12
if C == 4 + num_classes:
preds = data.T # (16, 8400) -> (8400, 16)
elif data.shape[1] == 4 + num_classes:
preds = data # 已经是 (8400, 16)
else:
print(f"[ERROR] 通道数 {C} 与期望 {4 + num_classes} 不符")
return [], [], []
print(f"[DEBUG] preds shape after transpose: {preds.shape}")
boxes_xywh = preds[:, :4] # (8400, 4)
cls_scores = preds[:, 4:] # (8400, 12)
# ============ 修改点 2：明确不做 sigmoid ============
# 通道统计已确认 ch4~ch15 是 0~1 的概率值，
# 再套 sigmoid 会让背景点变成 ~0.5，导致 8400 个框全过阈值
class_ids = np.argmax(cls_scores, axis=1)
confidences = np.max(cls_scores, axis=1)
# 置信度过滤（0.25和ONNX一致）
mask = confidences >= conf_thresh
boxes_xywh = boxes_xywh[mask]
confidences = confidences[mask]
class_ids = class_ids[mask]
print(f"[DEBUG] 置信度过滤后: {len(boxes_xywh)} 个框")
if len(boxes_xywh) == 0:
return [], [], []
# xywh (center-x, center-y, width, height) -> xyxy
cx = boxes_xywh[:, 0]
cy = boxes_xywh[:, 1]
bw = boxes_xywh[:, 2]
bh = boxes_xywh[:, 3]
x1 = cx - bw / 2
y1 = cy - bh / 2
x2 = cx + bw / 2
y2 = cy + bh / 2
# 逆letterbox：缩放到原始图片尺寸
x1 = (x1 - pad_w) / scale
y1 = (y1 - pad_h) / scale
x2 = (x2 - pad_w) / scale
y2 = (y2 - pad_h) / scale
# NMSBoxes 需要 [[x,y,w,h], ...] 的 list
bboxes_list = [[float(x1[i]), float(y1[i]), float(x2[i] - x1[i]), float(y2[i] - y1[i])] for i in range(len(x1))]
scores_list = [float(confidences[i]) for i in range(len(confidences))]
indices = cv2.dnn.NMSBoxes(bboxes_list, scores_list, conf_thresh, iou_thresh)
if isinstance(indices, tuple):
indices = indices[0]
if len(indices) == 0:
return [], [], []
final_boxes = []
final_scores = []
final_cls = []
for i in indices:
final_boxes.append([x1[i], y1[i], x2[i], y2[i]])
final_scores.append(float(confidences[i]))
final_cls.append(int(class_ids[i]))
print(f"[DEBUG] NMS后: {len(final_boxes)} 个框")
return final_boxes, final_scores, final_cls
def postprocess_letterbox(outputs, scale, pad_w, pad_h, conf_thresh=0.25, iou_thresh=0.45):
data = np.array(outputs)
if data.ndim == 3 and data.shape[0] == 1:
data = data[0]
if data.ndim != 2:
return [], [], []
C, N = data.shape
num_classes = 12
if C == 4 + num_classes:
preds = data.T
elif data.shape[1] == 4 + num_classes:
preds = data
else:
return [], [], []
boxes_xywh = preds[:, :4]
cls_scores = preds[:, 4:]
class_ids = np.argmax(cls_scores, axis=1)
confidences = np.max(cls_scores, axis=1)
mask = confidences >= conf_thresh
boxes_xywh = boxes_xywh[mask]
confidences = confidences[mask]
class_ids = class_ids[mask]
if len(boxes_xywh) == 0:
return [], [], []
cx, cy, bw, bh = (boxes_xywh[:, 0], boxes_xywh[:, 1],
boxes_xywh[:, 2], boxes_xywh[:, 3])
x1 = cx - bw / 2.0
y1 = cy - bh / 2.0
x2 = cx + bw / 2.0
y2 = cy + bh / 2.0
x1 = (x1 - pad_w) / scale
y1 = (y1 - pad_h) / scale
x2 = (x2 - pad_w) / scale
y2 = (y2 - pad_h) / scale
bboxes_list = [[float(x1[i]), float(y1[i]),
float(x2[i] - x1[i]), float(y2[i] - y1[i])]
for i in range(len(x1))]
scores_list = [float(confidences[i]) for i in range(len(confidences))]
indices = cv2.dnn.NMSBoxes(bboxes_list, scores_list, conf_thresh, iou_thresh)
if isinstance(indices, tuple):
indices = indices[0]
if len(indices) == 0:
return [], [], []
final_boxes = [[x1[i], y1[i], x2[i], y2[i]] for i in indices]
final_scores = [float(confidences[i]) for i in indices]
final_cls = [int(class_ids[i]) for i in indices]
return final_boxes, final_scores, final_cls
# ================= 主循环 =================
prev_time = time.time()
while True:
ret, frame = cap.read()
if not ret:
print("读取帧失败，尝试重连...")
import time as t
t.sleep(1)
cap.release()
cap = cv2.VideoCapture(RTSP_URL, cv2.CAP_FFMPEG)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
continue
h, w = frame.shape[:2]
# Letterbox预处理（和ONNX版本一致）
mat, scale, pad_w, pad_h = preprocess_letterbox(frame, INPUT_SIZE)
# ========== 纯推理计时起点 ==========
t0 = time.perf_counter()
# 推理 - 兼容不同ncnn版本的extractor API
ex = net.create_extractor()
# 尝试用name输入，不行就试index
try:
ex.input(input_names[0], mat)
except TypeError:
# 某些版本需要index而不是name
try:
ex.input(0, mat)
except Exception as e2:
print(f"[ERROR] input by index also failed: {e2}")
continue
code, out = ex.extract("out0")
infer_ms = (time.perf_counter() - t0) * 1000
# ========== 纯推理计时终点 ==========
if code != 0:
print(f"提取失败 code={code}")
continue
# Post-process
boxes, scores, cls_ids = postprocess_letterbox(out, scale, pad_w, pad_h, CONF_THRESH, NMS_THRESH)
# 画框
for box, score, cid in zip(boxes, scores, cls_ids):
x1, y1, x2, y2 = [int(v) for v in box]
# 边界保护
x1 = max(0, x1)
y1 = max(0, y1)
x2 = min(w - 1, x2)
y2 = min(h - 1, y2)
label = f"{NAMES[cid] if cid < len(NAMES) else 'cls'+str(cid)} {score:.2f}"
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(
frame, label, (x1, max(0, y1 - 6)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2
)
# FPS
now = time.time()
fps = 1.0 / (now - prev_time)
prev_time = now
cv2.putText(frame, f"FPS:{fps:.1f} | Infer:{infer_ms:.1f}ms",
(10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
cv2.imshow("RTSP NCNN", frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
cap.release()
cv2.destroyAllWindows()verify_onnx.py
python
import cv2
import numpy as np
import onnxruntime as ort
def preprocess(image, size=640):
h, w = image.shape[:2]
scale = size / max(h, w)
new_w = int(w * scale)
new_h = int(h * scale)
resized = cv2.resize(image, (new_w, new_h))
pad_w = (size - new_w) // 2
pad_h = (size - new_h) // 2
padded = cv2.copyMakeBorder(
resized,
pad_h, size - new_h - pad_h,
pad_w, size - new_w - pad_w,
cv2.BORDER_CONSTANT,
value=(114, 114, 114)
)
blob = padded[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) / 255.0
blob = np.expand_dims(blob, 0)
return blob, scale, pad_w, pad_h
def postprocess(outputs, scale, pad_w, pad_h, conf_thresh=0.25, iou_thresh=0.45):
preds = outputs[0][0].T # (8400, 4 + nc)
boxes_xywh = preds[:, :4]
scores = preds[:, 4:]
class_ids = np.argmax(scores, axis=1)
confidences = np.max(scores, axis=1)
mask = confidences >= conf_thresh
boxes_xywh = boxes_xywh[mask]
confidences = confidences[mask]
class_ids = class_ids[mask]
if len(boxes_xywh) == 0:
return [], [], []
# xywh -> xyxy
cx = boxes_xywh[:, 0]
cy = boxes_xywh[:, 1]
bw = boxes_xywh[:, 2]
bh = boxes_xywh[:, 3]
x1 = cx - bw / 2
y1 = cy - bh / 2
x2 = cx + bw / 2
y2 = cy + bh / 2
# 转成 Python list，否则 NMSBoxes 会报类型错
bboxes_list = [
[float(x1[i]), float(y1[i]), float(x2[i]), float(y2[i])]
for i in range(len(x1))
]
scores_list = [float(confidences[i]) for i in range(len(confidences))]
indices = cv2.dnn.NMSBoxes(
bboxes_list,
scores_list,
conf_thresh,
iou_thresh
)
if isinstance(indices, tuple):
indices = indices[0]
if len(indices) == 0:
return [], [], []
final_boxes = []
final_scores = []
final_cls = []
for i in indices:
x1_, y1_, x2_, y2_ = x1[i], y1[i], x2[i], y2[i]
x1_ = (x1_ - pad_w) / scale
y1_ = (y1_ - pad_h) / scale
x2_ = (x2_ - pad_w) / scale
y2_ = (y2_ - pad_h) / scale
final_boxes.append([x1_, y1_, x2_, y2_])
final_scores.append(float(confidences[i]))
final_cls.append(int(class_ids[i]))
return final_boxes, final_scores, final_cls
if __name__ == "__main__":
session = ort.InferenceSession(
"best.onnx",
providers=["CPUExecutionProvider"]
)
input_name = session.get_inputs()[0].name
image = cv2.imread("test2.jpg") # 换成你的验证集图片
blob, scale, pad_w, pad_h = preprocess(image)
outputs = session.run(None, {input_name: blob})
boxes, scores, cls_ids = postprocess(outputs, scale, pad_w, pad_h)
names = [
"target1", "target2", "target3", "target4",
"target5", "target6", "target7", "target8",
"target9", "target10", "target11", "target12"
]
for box, score, cid in zip(boxes, scores, cls_ids):
x1, y1, x2, y2 = [int(v) for v in box]
label = f"{names[cid]} {score:.2f}"
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, label, (x1, y1 - 6),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
# for box, score, cid in zip(boxes, scores, cls_ids):
# x1, y1, x2, y2 = [int(v) for v in box]
# print(f"class={cid}, conf={score:.3f}, box=[{x1},{y1},{x2},{y2}]")
# cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.imwrite("onnx_result_class.jpg", image)