介绍过 YOLO 的背景、起源等知识后,今天就将如何使用 YOLO 进行 目标检测(Object Detection)。安装就不说了,ultralytics 官网 提供多种安装方式可选择。

通常使用场景就下面几个步骤:

  1. 准备数据集(俗称图片“打标”):借助打标工具,对一批图片进行标注,构建训练数据集。
  2. 模型训练:提供打标好的图片数据集,训练 YOLO 模型文件。
  3. 模型预测:基于已训练好的模型文件,对图片或视频进行预测验证。

1. 准备数据集

YOLO 支持多种任务类型,不同任务对数据格式要求不同:

  • 目标检测(Object Detection):图像 + 边界框标注
  • 图像分类(Image Classification):图像 + 类别标签
  • 实例分割(Instance Segmentation):图像 + 多边形掩码标注
  • 姿态估计(Pose Estimation):图像 + 关键点坐标

以下以 目标检测 为例介绍数据集准备方法。

1.1. 数据集结构

YOLO 目标检测数据集通常按以下目录组织:

dataset/
│
├── images/                # 存放所有图片
│   ├── train/              # 训练集图片
│   │    ├── img001.jpg
│   │    ├── img002.jpg
│   │    └── ...
│   ├── val/                # 验证集图片
│   │    ├── img101.jpg
│   │    ├── img102.jpg
│   │    └── ...
│   └── test/               # (可选)测试集图片
│        ├── img201.jpg
│        └── ...
│
├── labels/                # 存放标注文件(与 images 对应)
│   ├── train/              # 训练集标注
│   │    ├── img001.txt
│   │    ├── img002.txt
│   │    └── ...
│   ├── val/                # 验证集标注
│   │    ├── img101.txt
│   │    ├── img102.txt
│   │    └── ...
│   └── test/               # (可选)测试集标注
│        ├── img201.txt
│        └── ...
│
└── data.yaml               # 数据集配置文件

1.1.1. 标注文件格式

YOLO 的标注文件为 .txt 格式,每一行表示一个目标:

<class_id> <x_center> <y_center> <width> <height>
  • class_id:类别 ID,从 0 开始
  • x_centery_center:目标边界框中心点坐标(相对于图像宽高,范围 0~1)
  • widthheight:边界框宽高(相对于图像宽高,范围 0~1)

例如:

0 0.512 0.423 0.134 0.276
1 0.325 0.600 0.250 0.400

1.1.2. 类别定义文件

data.yaml 文件用于定义数据集路径与类别名称,例如:

path: /xxx     #train/val/test的根目录 
train: dataset/images/train
val: dataset/images/val
test: dataset/images/test   # 可选

nc: 3
names: ['cat', 'dog', 'person']
  • trainval:训练集与验证集的路径
  • nc:类别数
  • names:类别名称列表

1.2. 划分数据集

  • train(训练集)
    用于模型参数更新,数据量应占多数,保证模型能学习到多样化的特征。
  • val(验证集)
    用于训练过程中评估模型效果(mAP、Precision、Recall),帮助调节超参数、防止过拟合。
  • test(测试集,可选)
    用于最终评估模型的泛化能力,不参与训练和验证过程。
图片数量比例推荐

在实际项目中,常见的划分比例有:

集合类型推荐比例说明
训练集 train70% ~ 80%占多数,保证模型学习到足够的特征
验证集 val10% ~ 20%用于训练过程中的效果评估
测试集 test10%(可选)最终评估模型泛化能力
示例:

如果你有 1000 张图片

  • 训练集:800 张
  • 验证集:150 张
  • 测试集:50 张(可选)
划分数据集的注意事项
  1. 随机划分
    保证不同集合中的数据分布一致,防止验证集/测试集与训练集分布差异过大。
  2. 类别均衡
    确保每个类别在 train / val / test 中都有出现,避免某些类别只出现在训练集或验证集中。
  3. 场景多样化
    各集合中应包含不同光照、角度、背景的样本,提升模型泛化能力。
  4. 文件名对应
    images/train/img001.jpg 必须对应 labels/train/img001.txt,文件名(不含后缀)一致。

1.3.数据标注工具

常用的标注工具有:LabelImg、Labelme、LabelStudio、Roboflow 等。

标注工具的功能都是:基于一张图片,人工在图片上标注出想要定义的区域,生成标注文件。

前面说了,YOLO 的标注文件为 .txt 格式,内容也有自己的格式要求。所以选择标注工具时要看是否天然支持 YOLO 的格式。

实际使用后推荐两种标注工具:

  • LabelImg:是一个轻量级、单一任务的开源标注工具,适合快速制作目标检测数据集,学习成本低,功能聚焦。
  • Label Studio:是一个通用的、可扩展的标注平台,适合需要多种数据类型标注、团队协作、与机器学习流程深度集成的场景。

个人快速使用推荐 LabelImg,团队项目使用还是得用 Label Studio。

实际上二者都是同一家公司的开源产品,Humansignal 是美国的一家公司,专注于数据标注与数据管理平台,既有开源工具,也有商业化 SaaS 平台。

两者定位对比

特性LabelImgLabel Studio
数据类型仅图像(目标检测)多模态(文本、图像、音频、视频、时间序列等)
功能范围单一任务,轻量级多任务,支持团队协作与质量控制
使用场景快速制作小型目标检测数据集构建完整标注流水线,适合中大型项目
部署方式本地运行本地部署 / 云端部署
用户群体开发者、研究人员企业团队、科研团队、AI产品开发者

1.3.1. LabelImg

上手难度最低,详细可参考 LabelImg Github

提供的是桌面客户端,界面简约易上手。但因为是客户端,所以对本地电脑上 python版本等环境有自己的一些要求。不要相信wiki里Docker的安装方式,mac上有很多问题。

下面附上一键安装、启动的脚本。

安装 install.sh
#!/bin/bash

echo "=== 检查 Homebrew 是否安装 ==="
if ! command -v brew &> /dev/null; then
    echo "Homebrew 未安装,开始安装..."
    /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
    echo 'eval "$(/opt/homebrew/bin/brew shellenv)"' >> ~/.zprofile
    eval "$(/opt/homebrew/bin/brew shellenv)"
else
    echo "Homebrew 已安装"
fi

echo "=== 安装 Homebrew Python ==="
brew install python

echo "=== 安装 Qt5 和 libxml2 ==="
brew install qt@5 libxml2
brew link --force qt@5

echo "=== 创建虚拟环境 ==="
python3 -m venv venv
source venv/bin/activate

echo "=== 安装 PyQt5 和 lxml ==="
pip install --upgrade pip
pip install pyqt5 lxml

echo "=== 下载 labelImg ==="
if [ ! -d "labelImg" ]; then
    git clone https://github.com/heartexlabs/labelImg.git
fi
cd labelImg

echo "=== 编译 labelImg (Qt5 + Python3) ==="
make qt5py3

echo "=== 启动 labelImg ==="
python3 labelImg.py

echo "=== 完成!你可以用以下命令启动 labelImg ==="
echo "source venv/bin/activate && cd labelImg && python3 labelImg.py"
启动:start.sh
#!/bin/bash
# 启动 labelImg 脚本

# 切换到当前脚本所在目录(保证路径正确)
cd "$(dirname "$0")" || { echo "无法进入脚本所在目录"; exit 1; }

# 激活虚拟环境
source venv/bin/activate

# 进入 labelImg 目录
cd labelImg || { echo "找不到 labelImg 目录"; exit 1; }

# 如果有参数就传给 labelImg.py(图片路径 / 类别文件)
if [ $# -eq 0 ]; then
    python3 labelImg.py
else
    python3 labelImg.py "$@"
fi

1.3.2. LabelStudio

LabelStudio 的安装使用可以看 LabelStudio 官网WIKI

它的功能丰富的多,而且支持团队协作,在使用时上手难度会高一点。但因为提供的 WEB 应用,安装非常容易,下面提供 Docker 一键安装命令。

#!/bin/bash
# 文件名: start_label_studio.sh
# 用法: ./start_label_studio.sh

# 本机端口
HOST_PORT=8090

# 数据目录(持久化账号和项目数据)
DATA_DIR="$(pwd)/my_label_studio"

# 默认账号密码
DEFAULT_USERNAME="admin@example.com"
DEFAULT_PASSWORD="123456"

# 容器名称
CONTAINER_NAME="label_studio"

# 创建数据目录(如果不存在)
mkdir -p "$DATA_DIR"

echo "后台启动 Label Studio..."
echo "本机端口: $HOST_PORT"
echo "数据目录: $DATA_DIR"
echo "默认账号: $DEFAULT_USERNAME"
echo "默认密码: $DEFAULT_PASSWORD"
echo "容器名称: $CONTAINER_NAME"

# 如果容器已存在,先停止并删除
if [ "$(docker ps -aq -f name=${CONTAINER_NAME})" ]; then
    echo "已有容器存在,先停止并删除..."
    docker stop ${CONTAINER_NAME} >/dev/null
    docker rm ${CONTAINER_NAME} >/dev/null
fi

# 后台启动 Docker 容器
docker run -d \
    --name ${CONTAINER_NAME} \
    -p ${HOST_PORT}:8080 \
    -v "${DATA_DIR}:/label-studio/data" \
    -e LABEL_STUDIO_USERNAME="${DEFAULT_USERNAME}" \
    -e LABEL_STUDIO_PASSWORD="${DEFAULT_PASSWORD}" \
    heartexlabs/label-studio:latest

echo "启动完成!访问: http://localhost:${HOST_PORT}"

2. 模型训练

模型训练示例命令:

yolo train model=yolov11n.pt data=data.yaml epochs=100 imgsz=640 batch=16
  • model=yolov11n.pt → 选择小型网络(较少卷积层,推理快)
  • epochs=100 → 让网络权重在训练集上迭代 100 次
  • imgsz=640 → 输入图片缩放到 640×640,保持足够细节
  • batch=16 → 每次用 16 张图片更新一次权重

总结(神经网络视角)

  • model → 网络结构规模(神经元数量、卷积层宽度/深度)
  • epochs → 学习轮次(权重更新次数)
  • imgsz → 输入特征的分辨率(影响卷积特征图细节)
  • batch → 每次梯度更新的样本数(影响梯度稳定性与显存占用)

2.1. 参数说明

2.1.1. model

  • 作用:选择神经网络的结构和预训练权重。
  • 神经网络背景

    • YOLOv11 的不同版本(n/s/m/l/x)本质上是同一架构的不同规模,区别在于:

      • 卷积层的数量(深度)
      • 每层的通道数(宽度)
    • 小模型(如 yolov11n)参数量少,计算量低,推理速度快,但表达能力有限。
    • 大模型(如 yolov11x)参数量多,可以拟合更复杂的数据分布,但训练时间长、显存占用高。
  • 类比

    • 模型大小就像大脑的神经元数量:更多神经元(大模型)有更强的学习能力,但需要更多训练时间和能量(显存/算力)。
  • 建议

    • 数据少/硬件弱 → 小模型
    • 数据多/硬件强 → 大模型

2.1.2. epochs

  • 作用:训练集被完整遍历的次数。
  • 神经网络背景

    • 每次遍历(一个 epoch),模型会在所有样本上更新参数一次。
    • 在反向传播(Backpropagation)中,梯度会逐步调整权重,让网络更好拟合数据。
    • 如果 epochs 太少 → 权重还没收敛,模型欠拟合。
    • 如果 epochs 太多 → 网络可能记住训练集细节,导致过拟合(泛化能力下降)。
  • 类比

    • Epochs 就像复习次数:复习太少记不住,复习太多可能只会记住试卷答案。
  • 建议

    • 小数据集:更多 epochs(100~300)
    • 大数据集:适中 epochs(50~150),并结合早停(early stopping)

2.1.3. imgsz

  • 作用:训练输入图片的分辨率。
  • 神经网络背景

    • 在卷积神经网络(CNN)中,输入尺寸决定了特征图(feature map)的大小。
    • 高分辨率 → 保留更多细节(尤其对小目标检测有利),但卷积计算量和显存占用增加。
    • 低分辨率 → 特征图更小,计算快,但可能丢失细节信息。
    • YOLOv11 会将所有输入图片缩放到 imgsz × imgsz,保证输入统一。
  • 类比

    • imgsz 就像照片的像素数:更清晰的照片可以看到小物体,但处理起来更慢。
  • 建议

    • GPU 显存 ≥ 8GB → 640 或更高
    • 显存较小 → 512 或 416

2.1.4. batch

  • 作用:每次梯度更新时使用的样本数量。
  • 神经网络背景

    • 在训练中,数据不是一次性全部送入网络,而是分成小批次(batch)。
    • 批量大小影响:

      • 梯度估计的稳定性:大 batch → 梯度更平滑,训练更稳定;小 batch → 梯度波动大,但泛化能力可能更好。
      • 显存占用:batch 越大,显存需求越高。
    • YOLOv11 的训练过程:每个 batch 输入网络 → 前向传播(forward) → 计算损失(loss) → 反向传播(backward) → 更新权重。
  • 类比

    • Batch 就像一次课堂的学生人数:人数多(大 batch)统计更稳定,但需要更大教室(显存)。
  • 建议

    • 显存 4GB:batch ≤ 8
    • 显存 8GB:batch 16~32
    • 显存 16GB:batch 32~64
    • 不确定时可用 batch=auto 自动适配

2.2. 预训练权重

在模型训练时,都需要选择一个预训练模型,如:model=yolov11n.pt
yolov11n.pt 本身就是一个训练好的通用模型文件,可以直接用来预测常见场景图片。

那么如果我本地训练集里的图片,就包含2个分类。那么基于 yolov11n.pt这个模型,再训练出来的新模型,是只有2个分类,还是再原分类基础上再加2个分类(模型微调)?

答案是前者,只有 2个分类。

1. model 是什么?

在 YOLO 中,model 参数用来指定:

  • 网络结构定义(来自 .yaml 文件)
  • 以及可选的预训练权重(来自 .pt 文件)

例如:

model=yolov11n.pt       # 使用 nano 模型的 COCO 预训练权重
model=yolov11m.yaml     # 使用 medium 结构,从零开始训练
model=path/to/custom.pt # 用自己训练好的权重继续训练
2. yolov11n.pt 里到底有什么
  • yolov11n.pt 是 Ultralytics 提供的预训练权重,通常是用 COCO 数据集(80 类) 训练出来的。
  • 它包含两部分内容:

    1. 网络结构定义(模型的层、通道数等)
    2. 权重参数(卷积核、BN 参数等)
    3. 检测头的配置(最后一层分类数 = 80)
3. .yaml vs .pt 的区别
  • .yaml 文件:只包含模型结构定义(层数、通道数、模块类型等),不包含已训练好的权重
  • .pt 文件:包含模型结构 + 已训练好的权重(通常来自大规模数据集,比如 COCO)。
4. 如果用 yolov11n.pt 训练自己的 2 类数据,会怎样?

当你在训练命令中指定:

yolo detect train data=data.yaml model=yolov11n.pt

data.yaml 内容是:

names:
  0: cat
  1: dog

训练脚本会做两件事:

  1. 读取 data.yaml → 确认你只有 2 个类别。
  2. 自动修改检测头

    • 原本的输出层是 num_classes=80(COCO)
    • 会被替换成 num_classes=2
    • 原来的检测头权重会丢弃(因为输出维度不一样)
  3. 保留主干网络(Backbone)和颈部(Neck)的权重

    • 这些部分保留了在 COCO 上学到的通用特征(边缘、纹理、形状等)
    • 这就是 迁移学习 的核心:用大数据集学到的特征来加速小数据集的训练

最终结果:你训练出来的模型 只包含 2 类,不会混入原来 COCO 的 80 类。

5. 为什么要用 .pt 而不是 .yaml
  • 如果你数据量很小(几十张~几千张),从 .yaml 结构开始训练,相当于从零学特征,收敛慢、精度低。
  • .pt 预训练权重,前面的大部分网络参数已经学会了“看图”的能力,只需要学会区分你自己的类别即可。
6. 迁移学习的好处
  • 虽然最后一层分类器是新建的,但前面的特征提取部分保留了 COCO 上学到的特征。
  • 这些特征对很多常见物体(边缘、纹理、形状等)都有泛化能力,即使你的数据集只有 2 类,也能更快收敛、效果更好。
7. 如果你想保留原来的 80 类再加新类?
  • 这种需求叫 增量学习(Incremental Learning),YOLOv11 默认不直接支持。
  • 常规做法是:

    1. 在数据集中包含原有 80 类数据 + 新类数据。
    2. 修改 data.yaml,把 names 列表改成 81 类。
    3. 用原始权重初始化,并重新训练(可能需要调低学习率,防止遗忘)。
8. 结合神经网络原理理解
  • YOLOv11 模型分为:

    1. Backbone(特征提取)
    2. Neck(特征融合)
    3. Head(检测输出)
  • .pt 预训练权重的 Backbone + Neck 部分是通用的视觉特征提取器,相当于已经学会“看图”。
  • Head 部分是任务相关的分类器 + 回归器:

    • 类别数不同 → 必须重新初始化
    • 边框回归部分可以复用,因为它和类别数无关

3. 训练/验证/测试与数据

命令类型train 参与val 参与test 参与作用说明
yolo detect train✅ 权重更新✅ 验证指标训练+验证
yolo detect val✅ 验证指标单独验证
yolo detect val --split test✅ 测试指标最终测试
yolo detect predict✅(或任意路径)推理输出
yolo export模型导出

3.1 训练

yolo detect train data=data.yaml model=yolov8n.pt epochs=100 imgsz=640

数据参与情况:

  • images/train参与训练阶段(权重更新)
  • images/val参与验证阶段(每个 epoch 结束后评估性能)
  • images/test不参与(训练命令不会用测试集)

流程说明:

  1. 训练阶段

    • images/train(及对应 labels/train)中按 batch 读取数据
    • 前向传播 → 计算损失 → 反向传播 → 更新权重
  2. 验证阶段

    • images/val(及对应 labels/val)读取数据
    • 计算 mAP、Precision、Recall 等指标
    • 不更新权重
  3. 保存模型

    • 根据验证集表现保存 best.pt

3.2 验证命令

yolo detect val data=data.yaml model=best.pt imgsz=640

数据参与情况:

  • images/train不参与
  • images/val参与验证
  • images/test不参与(除非你修改 data.yamlval 指向测试集路径)

流程说明:

  • 读取 images/val,用指定模型(best.pt)进行推理
  • 计算验证集上的指标(mAP、Precision、Recall)
  • 常用于单独评估模型在验证集上的表现

3.3 测试命令

yolo detect val data=data.yaml split=test model=best.pt imgsz=640

或者:

yolo detect val data=data.yaml model=best.pt imgsz=640 --split test

数据参与情况:

  • images/train不参与
  • images/val不参与
  • images/test参与测试

流程说明:

  • 读取 images/test,用指定模型进行推理
  • 计算测试集上的指标
  • 常用于最终评估模型泛化能力

3.4 推理命令

yolo detect predict model=best.pt source=dataset/images/test

数据参与情况:

  • images/train不参与
  • images/val不参与
  • images/test(或任意路径)→ 参与推理

流程说明:

  • 读取 source 指定路径的图片
  • 用模型进行推理
  • 输出预测结果(带边框图片、标签文件)

3.5 导出命令

yolo export model=best.pt format=onnx

数据参与情况:

  • 所有数据集目录 → 不参与
  • 导出模型到指定格式(ONNX、TensorRT 等)
  • 与数据无关

4. 模型推理

4.1. 推理流程

YOLO在推理预测阶段的核心任务是:

将输入图像快速、准确地检测出目标的位置(边界框)、类别以及置信度分数

目标检测的推理过程如下。

4.1.1. 输入处理

  • 读取图像:可以是本地文件、视频帧、摄像头流等。
  • 缩放与填充:YOLO 默认将输入缩放到模型的固定尺寸(例如 640×640),使用 letterbox 填充保持比例。
  • 颜色通道调整:通常将 BGR(OpenCV 默认)转为 RGB。
  • 归一化:像素值从 [0, 255] 转为 [0, 1] 浮点数。
  • 维度变换:形状由 (H, W, 3) 转为 (3, H, W),再添加 batch 维度。
在 Ultralytics YOLO 的 Python API 中,这些步骤会自动完成。

4.1.2. 模型前向推理

  • 将预处理后的图像张量输入 YOLOv11 模型。
  • 模型内部结构:

    • Backbone(主干网络):提取图像特征(CSPDarknet、改进的 Conv 模块等)。
    • Neck(特征融合):如 FPN+PAN 结构,将不同尺度的特征融合,便于检测不同大小的目标。
    • Head(检测头):输出预测结果,包括:

      • 边界框参数(中心点 x,y,宽 w,高 h)
      • 类别概率分布
      • 置信度分数

输出通常是一个形状为 (batch, num_preds, 4 + num_classes) 的张量。

4.1.3. 后处理

推理输出的原始张量需要进一步处理才能得到最终的检测结果。

边界框解码
  • 模型输出的坐标是相对于特征图的,需要通过公式映射回原图尺寸。
  • YOLO 使用了 Anchor-free 设计,预测的是相对于网格单元的偏移量。
置信度计算
  • 置信度 = 目标存在概率 × 类别概率。
  • 过滤低置信度的预测(例如 conf < 0.25)。
NMS(非极大值抑制)
  • 解决多个框重复检测同一目标的问题。
  • 保留置信度最高的框,去掉与其 IoU(交并比)超过阈值的其他框。
  • YOLO 默认使用 加权 NMS标准 NMS,可选 Soft-NMS。

4.1.4. 输出结果

最终返回:

  • 边界框坐标(在原图上的位置)
  • 类别 ID / 名称
  • 置信度分数

4.2. API服务

因为都是 python 环境,所以和 OCR 的API服务端实现方式一样:FastAPI + uvicorn

下面是提供一个API的 python代码:

  • 预加载模型文件进行预测
  • 基于每次上传的模型文件进行预测(用于测试)

    • 可测试返回JSON
    • 可测试返回标注后的图片
    • 可测试返回标注后的视频(浏览器可直接播放)

Python代码:

from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse, StreamingResponse
from ultralytics import YOLO
from PIL import Image
from pathlib import Path
import io
import os
import cv2
import tempfile
import subprocess
import requests
import shutil


app = FastAPI()

# 缓存:模型路径 & YOLO对象
model_path_cache = {}  # key: version_code, value: model_path
model_obj_cache = {}   # key: version_code, value: YOLO object

# 模型存放目录
MODEL_DIR = Path("/app/yolo/models")
MODEL_DIR.mkdir(parents=True, exist_ok=True)

# 管理域名
MNG_DOMAIN = os.getenv("MNG_DOMAIN", "http:xxx")


def get_latest_model_info(code: str) -> dict:
    """调用获取最新模型API"""
    url = f"{MNG_DOMAIN}/{api_path}?code={code}"
    resp = requests.get(url)
    resp.raise_for_status()
    data = resp.json()
    if not data.get("success"):
        raise ValueError(f"获取最新模型失败: {data.get('errorMessage')}")
    return data["data"]


def download_model_file(url: str, save_path: Path):
    """下载模型文件"""
    resp = requests.get(url, stream=True)
    resp.raise_for_status()
    with open(save_path, "wb") as f:
        shutil.copyfileobj(resp.raw, f)


def load_image_from_upload(upload_file: UploadFile) -> Image.Image:
    """加载上传的图片"""
    image_bytes = upload_file.file.read()
    return Image.open(io.BytesIO(image_bytes)).convert("RGB")


def parse_results(model: YOLO, results) -> list:
    """解析 YOLO 预测结果"""
    detections = []
    for r in results:
        for box in r.boxes:
            cls_id = int(box.cls[0])
            score = float(box.conf[0])
            xyxy = box.xyxy[0].tolist()
            detections.append({
                "class_id": cls_id,
                "class_name": model.names[cls_id],
                "confidence": score,
                "bbox": xyxy
            })
    return detections


def cleanup_old_model(code: str):
    """删除旧版本模型及缓存"""
    keys_to_remove = [k for k in list(model_path_cache.keys()) if k.endswith(f"_{code}")]
    for k in keys_to_remove:
        old_path = model_path_cache[k]
        if os.path.exists(old_path):
            os.remove(old_path)
        model_path_cache.pop(k, None)
        model_obj_cache.pop(k, None)


@app.post("/predict")
async def predict(
        code: str = Form(...),
        version: str = Form(...),
        file: UploadFile = File(...)
):
    key = f"{version}_{code}"

    # 如果缓存中已有该版本模型
    if key in model_obj_cache:
        model = model_obj_cache[key]
        used_version = version
    else:
        # 获取最新模型信息
        latest_info = get_latest_model_info(code)
        latest_version = latest_info["version"]
        latest_url = f"{MNG_DOMAIN}{latest_info['url']}"

        if latest_version == version:
            # 请求版本是最新版本 -> 下载并缓存
            cleanup_old_model(code)
            model_path = MODEL_DIR / f"{version}_{code}.pt"
            download_model_file(latest_url, model_path)
            model_path_cache[key] = str(model_path)
            model_obj_cache[key] = YOLO(str(model_path))
            model = model_obj_cache[key]
            used_version = version
        else:
            # 最新版本与请求版本不一致
            # 在缓存中找该code的任意版本
            matched_keys = [k for k in model_obj_cache.keys() if k.endswith(f"_{code}")]
            if matched_keys:
                # 用缓存中的第一个版本
                cache_key = matched_keys[0]
                model = model_obj_cache[cache_key]
                used_version = cache_key.split("_")[0]
            else:
                # 缓存中没有该code -> 下载最新版本
                cleanup_old_model(code)
                new_key = f"{latest_version}_{code}"
                model_path = MODEL_DIR / f"{latest_version}_{code}.pt"
                download_model_file(latest_url, model_path)
                model_path_cache[new_key] = str(model_path)
                model_obj_cache[new_key] = YOLO(str(model_path))
                model = model_obj_cache[new_key]
                used_version = latest_version

    # 预测
    img = load_image_from_upload(file)
    results = model.predict(img)
    detections = parse_results(model, results)

    return JSONResponse(content={"detections": detections})




################ 测试接口 ################


@app.post("/test")
async def test(model_file: UploadFile = File(...), image_file: UploadFile = File(...)):
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as tmp_model:
            tmp_model.write(await model_file.read())
            tmp_model_path = tmp_model.name

        model = YOLO(tmp_model_path)
        os.remove(tmp_model_path)

        img = load_image_from_upload(image_file)
        results = model.predict(img)
        detections = parse_results(model, results)

        return JSONResponse(content={"detections": detections})

    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)


@app.post("/test/image")
async def test_image(model_file: UploadFile = File(...), image_file: UploadFile = File(...)):
    try:
        # 保存临时模型
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as tmp_model:
            tmp_model.write(await model_file.read())
            tmp_model_path = tmp_model.name

        model = YOLO(tmp_model_path)
        os.remove(tmp_model_path)

        # 保存临时图片
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_img:
            tmp_img.write(await image_file.read())
            tmp_img_path = tmp_img.name

        # 推理并保存标注图片
        results = model.predict(tmp_img_path, save=True, project="/tmp", name="yolo_output", exist_ok=True)

        # 修复:转换为 Path 再拼接
        output_path = Path(results[0].save_dir) / os.path.basename(tmp_img_path)

        with open(output_path, "rb") as f:
            img_bytes = f.read()

        os.remove(tmp_img_path)

        return StreamingResponse(io.BytesIO(img_bytes), media_type="image/jpeg")

    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)



MAX_VIDEO_SIZE_MB = 5  # 最大视频大小(MB)
BATCH_SIZE = 8         # 批量推理帧数

@app.post("/test/video")
async def test_video(model_file: UploadFile = File(...), video_file: UploadFile = File(...)):
    tmp_files = []  # 记录临时文件,方便统一删除
    try:
        # 检查视频大小
        video_file.file.seek(0, os.SEEK_END)
        size_mb = video_file.file.tell() / (1024 * 1024)
        video_file.file.seek(0)
        if size_mb > MAX_VIDEO_SIZE_MB:
            return JSONResponse(
                content={"error": f"视频文件过大 ({size_mb:.2f} MB),最大允许 {MAX_VIDEO_SIZE_MB} MB"},
                status_code=400
            )

        # 保存模型文件
        tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pt")
        tmp_files.append(tmp_model.name)
        tmp_model.write(await model_file.read())
        tmp_model.close()

        # 加载模型(自动使用 GPU,如果可用)
        model = YOLO(tmp_model.name)
        device = 0 if model.device.type != "cpu" else "cpu"

        # 保存输入视频
        tmp_input = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
        tmp_files.append(tmp_input.name)
        tmp_input.write(await video_file.read())
        tmp_input.close()

        # 打开视频
        cap = cv2.VideoCapture(tmp_input.name)
        if not cap.isOpened():
            return JSONResponse(content={"error": "无法打开视频"}, status_code=400)

        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # 中间视频文件(MJPG)
        tmp_mid = tempfile.NamedTemporaryFile(delete=False, suffix=".avi")
        tmp_files.append(tmp_mid.name)
        out = cv2.VideoWriter(tmp_mid.name, cv2.VideoWriter_fourcc(*"MJPG"), fps, (width, height))

        # 批量推理
        frames_batch = []
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frames_batch.append(frame)

            if len(frames_batch) == BATCH_SIZE:
                results = model.predict(frames_batch, device=device, verbose=False)
                for res in results:
                    out.write(res.plot())
                frames_batch.clear()

        # 处理剩余的帧
        if frames_batch:
            results = model.predict(frames_batch, device=device, verbose=False)
            for res in results:
                out.write(res.plot())

        cap.release()
        out.release()

        # 转码为浏览器可播放的 MP4(H.264 Baseline)
        tmp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
        tmp_files.append(tmp_output.name)
        subprocess.run([
            "ffmpeg", "-y", "-i", tmp_mid.name,
            "-c:v", "libx264", "-preset", "fast", "-profile:v", "baseline",
            "-level", "3.0", "-pix_fmt", "yuv420p",
            "-movflags", "+faststart",
            tmp_output.name
        ], check=True)

        # 返回视频流
        video_stream = open(tmp_output.name, "rb")
        return StreamingResponse(video_stream, media_type="video/mp4")

    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)
    finally:
        # 确保删除所有临时文件
        for f in tmp_files:
            try:
                os.remove(f)
            except FileNotFoundError:
                pass

Dockerfile:

FROM ultralytics/ultralytics:latest
WORKDIR /ultralytics/workspace

COPY app.py .
COPY xxx.pt ./xxx.pt

RUN pip install --no-cache-dir fastapi uvicorn python-multipart Pillow

# 暴露端口
EXPOSE 9003

# 默认启动命令
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "9003", "--workers", "3"]

还有个 rapidocr + yolo 打包构建的 Dockerfile

FROM ultralytics/ultralytics:latest

ENV DEBIAN_FRONTEND=noninteractive

WORKDIR /app

# 安装系统工具
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        vim \
        ffmpeg \
        curl \
        iputils-ping \
        net-tools \
        dnsutils \
        inetutils-traceroute \
        telnet \
        procps && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# 在 Conda Python 环境中安装 Python 依赖
RUN /opt/conda/bin/pip install --no-cache-dir \
    requests \
    fastapi \
    uvicorn[standard] \
    python-multipart \
    pillow \
    numpy \
    rapidocr \
    onnxruntime \
    opencv-python \
    -i https://mirrors.aliyun.com/pypi/simple

# 拷贝 OCR 代码
WORKDIR /app/rapidocr
COPY ocr.py /app/rapidocr/ocr.py

# 拷贝 YOLO 代码和模型
WORKDIR /app/yolo
COPY yolo.py /app/yolo/yolo.py
COPY xxx.pt /app/yolo/xxx.pt

# 回到主目录
WORKDIR /app

# 暴露端口
EXPOSE 9000 9001

# 启动两个服务
CMD bash -c "cd /app/rapidocr && /opt/conda/bin/uvicorn ocr:app --host 0.0.0.0 --port 9000 --workers 6 & \
             cd /app/yolo && /opt/conda/bin/uvicorn yolo:app --host 0.0.0.0 --port 9001 --workers 3 & \
             wait"

5. 训练建议

如果需要模型训练,建议安装环境有GPU可用,基于GPU训练,而不是仅通过CPU。

本地试过通过 Docker 安装容器,用 CPU 训练后,预加载权重文件用 n 勉强还行,macbook pro 风扇转的飞起。用 s 就直接跑不了。

好在 yolo 官网支持 Apple Silicon MPS 训练,使用 Apple 芯片的GPU进行训练。干脆直接在 mac上安装,而不用 Docker训练了。

mac 安装脚本 install.sh

#!/bin/bash

# ============================================
# 安装 Python 3.11.12 + ultralytics (YOLOv8)
# 适用于 Mac M4 芯片 (Apple Silicon ARM64)
# 作者: ChatGPT
# ============================================

set -e

echo "==== 检查 pyenv 是否安装 ===="
if ! command -v pyenv >/dev/null 2>&1; then
    echo "pyenv 未安装,开始安装..."
    # 如果 brew 可用,用 brew 安装
    if command -v brew >/dev/null 2>&1; then
        brew install pyenv
    else
        echo "Homebrew 不可用,直接从源码安装 pyenv..."
        git clone https://github.com/pyenv/pyenv.git ~/.pyenv
    fi
else
    echo "pyenv 已安装"
fi

# 配置 pyenv 环境变量
echo "==== 配置 pyenv 环境变量 ===="
if ! grep -q 'pyenv init' ~/.zshrc; then
cat >> ~/.zshrc <<'EOF'
export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init --path)"
eval "$(pyenv init -)"
EOF
fi
export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init --path)"
eval "$(pyenv init -)"

# 安装 Python 3.11.12
echo "==== 安装 Python 3.11.12 ===="
if ! pyenv versions | grep -q "3.11.12"; then
    pyenv install 3.11.12
else
    echo "Python 3.11.12 已安装"
fi

# 设置全局 Python 版本
pyenv global 3.11.12
echo "当前 Python 版本: $(python --version)"

# 升级 pip
echo "==== 升级 pip/setuptools/wheel ===="
pip install -U pip setuptools wheel

# 安装 ultralytics
echo "==== 安装 ultralytics (YOLOv8) ===="
pip install -U ultralytics

# 测试 YOLO
echo "==== 测试 YOLO ===="
if command -v yolo >/dev/null 2>&1; then
    yolo
else
    echo "YOLO安装失败,请检查 pip 输出日志。"
fi

echo "==== 安装完成 ===="

文件结构(简述)

├── datasets
│   └── base_component
│       ├── dataset.yaml
│       ├── images
│       │   ├── train
│       │   │   ├── 20251208112810_74_101.jpg
│       │   │   ├── 20251208112811_75_101.jpg
│       │   │   ├── 20251208112812_76_101.jpg
│       │   └── val
│       │       ├── 20251208115935_98_101.jpg
│       └── labels
│           ├── train
│           │   ├── 20251208112810_74_101.txt
│           │   ├── 20251208112811_75_101.txt
│           │   ├── 20251208112812_76_101.txt
│           ├── train.cache
│           ├── val
│           │   ├── 20251208115935_98_101.txt
│           └── val.cache

训练命令:

yolo detect train data=datasets/base_component/dataset.yaml model=yolo11n.pt epochs=240 imgsz=896 batch=8 device=mps plots=True

数据集里40多张图片,训练 2个类别。未提高精度 imgsz=896,但带来就是显存占用提高,因此将 batch 从默认值 16 降到 8。
当前参数在 macbook pro m4 芯片上已是极限,GPU 占用在 93%左右。

训练结束后会在当前目录生成 runs 目录,包含每次训练的模型,找 best.pt(表现最好的权重)。

因为 plots=Trueruns/detect/train 目录中还会包含训练的可视化图片。

kerry@wuchenruideMacBook-Pro train % ls
args.yaml            confusion_matrix_normalized.png    train_batch0.jpg        train_batch2.jpg
BoxF1_curve.png            confusion_matrix.png        train_batch1.jpg        val_batch0_labels.jpg
BoxP_curve.png            labels.jpg            train_batch180.jpg        val_batch0_pred.jpg
BoxPR_curve.png            results.csv            train_batch181.jpg        weights
BoxR_curve.png            results.png            train_batch182.jpg

6. 训练流程

YOLO 模型训练时每个批次都会更新权重参数(除非用了梯度累积)。

  • 但是权重文件不会每个批次保存,只会在设定的周期(如每个 epoch 或最佳验证结果)保存。
  • 验证集只评估,不参与权重更新。

1. 训练过程的基本单位:批次(Batch)

  • 数据集被分成多个 批次(Batch),每个批次包含 batch_size 张图片。
  • 在 YOLO 训练中,每个批次会:

    1. 前向传播(Forward Pass):模型用当前权重预测这一批数据的输出(边界框、类别等)。
    2. 计算损失(Loss):基于预测结果和真实标签,计算总损失(box_loss、class_loss、dfl_loss 等)。
    3. 反向传播(Backward Pass):计算损失对权重的梯度。
    4. 更新权重(Optimizer Step):使用优化器(SGD、Adam 等)更新模型参数。

结论
每个批次结束时,模型权重都会更新一次(基于该批次的梯度)。


2. 权重更新的方式

  • 梯度累积(Gradient Accumulation)

    • 如果显存不足,可能会使用梯度累积,比如每 N 个批次再更新一次权重。
    • YOLO 在默认配置中通常是 每个批次都更新一次,除非你设置了 accumulate
  • 优化器的作用

    • 例如 Adam 会根据历史梯度信息调整学习率和更新方向。
    • SGD 则直接按梯度方向更新。

3. 权重文件保存的频率

  • 虽然权重每个批次都会更新,但 不会每个批次都保存权重文件,否则文件太多。
  • YOLO 的训练脚本通常会:

    • 每个 epoch(完整遍历一次训练集)结束后保存一次权重文件。
    • 或者只在验证集表现达到最佳时保存 best.pt
    • 还有一个 last.pt,表示最后一次训练结束时的权重。

4. 批次更新与验证集

  • 权重更新完全基于训练集的批次数据。
  • 验证集不会参与梯度更新,只用于评估当前权重的泛化能力。
  • 在每个 epoch 结束时,模型会用当前权重跑一遍验证集,计算 val/box_lossmAP 等指标。

5. 简单流程示意

假设一个 epoch 有 100 个 batch:

for epoch in range(num_epochs):
    for batch in train_loader:
        outputs = model(batch)             # 前向传播
        loss = loss_fn(outputs, labels)    # 计算损失
        loss.backward()                    # 反向传播
        optimizer.step()                   # 更新权重
        optimizer.zero_grad()              # 清空梯度
    validate(model, val_loader)            # 用验证集评估
    save_weights_if_needed(model)          # 保存权重文件

6. 流程图

下面是YOLO 训练流程时间线示意图,可以更直观理解:

  • 批次更新权重
  • 每个 epoch 结束验证
  • 保存权重文件

YOLO 训练流程(示意)

[Epoch 1]
 ├── Batch 1:  前向传播 → 计算损失 → 反向传播 → 更新权重
 ├── Batch 2:  前向传播 → 计算损失 → 反向传播 → 更新权重
 ├── Batch 3:  ...
 ├── ...
 ├── Batch N:  前向传播 → 计算损失 → 反向传播 → 更新权重
 └── 验证集评估 (val/box_loss, val/class_loss, mAP...)
     └── 保存权重文件(如 last.pt,若表现最好则保存 best.pt)

[Epoch 2]
 ├── Batch 1:  前向传播 → 计算损失 → 反向传播 → 更新权重
 ├── Batch 2:  ...
 ├── ...
 └── 验证集评估
     └── 保存权重文件(视情况保存)

[Epoch 3]
 ├── ...
 └── 验证集评估
     └── 保存权重文件

...

流程逻辑图(更直观)

┌─────────────┐
│ 读取训练批次 │
└──────┬──────┘
       │
       ▼
┌─────────────┐
│ 前向传播     │
└──────┬──────┘
       │
       ▼
┌─────────────┐
│ 计算损失     │
└──────┬──────┘
       │
       ▼
┌─────────────┐
│ 反向传播     │
└──────┬──────┘
       │
       ▼
┌─────────────┐
│ 更新权重     │ ← 每个批次都会发生
└──────┬──────┘
       │
       ▼
┌───────────────────────────┐
│ Epoch结束 → 用验证集评估模型 │
└───────────┬───────────────┘
            │
            ▼
┌───────────────────────────┐
│ 保存权重文件(last.pt / best.pt) │
└───────────────────────────┘
  • 每个 Batch:都会更新一次权重参数。
  • 每个 Epoch 结束:跑一次验证集,计算 val 损失和 mAP。
  • 保存权重文件:通常在每个 Epoch 或验证集表现最佳时保存。

KerryWu
679 声望171 粉丝

保持饥饿


引用和评论

0 条评论