YOLO模型训练(labelstudio支持多人标注)

图片标注及数据集训练模型

labelstud图片标注导出数据集

https://labelstud.io/

# Install the package
# into python virtual environment
pip install -U label-studio
# Launch it!
label-studio

#创建密码,并通过端口登录
label-studio start --username admin@example.com --password 123456

设置标注类型

labelstudio导出的压缩包目录如下

数据集

数据集文件夹上级目录下新建data.yaml,并根据文件夹层级编写data.yaml内容

手动划分数据集

labelstudio导出的数据集,手动将 images 和 labels 里的文件按 8:2 或 9:1 的比例，分别移动到 train 和 val 文件夹中

pole_dataset/
├── data.yaml          # 数据集配置文件,需要自己新建
├── train/
│   ├── images/
│   │   ├── img001.jpg
│   │   └── ...
│   └── labels/
│       ├── img001.txt  # YOLO 格式标注
│       └── ...
├── val/    #手动将 images 和 labels 里的文件按 8:2 或 9:1 的比例，分别移动到 train 和 val 文件夹中
│   ├── images/
│   │   └── ...
│   └── labels/
│       └── ...
└── test/
    ├── images/
    │   └── ...
    └── labels/
        └── ...

data.yaml 内容：

# 灯杆检测数据集配置
path: ./pole_dataset
train: train/images
val: val/images
test: test/images

# 类别
nc: 1  # 类别数量
names: ['pole']  # 类别名称,编号默认从0开始

不想划分数据集（直接训练）

my_dataset/          ← 数据集父目录
  ├── images/        ← 里面全是图片（001.jpg, 002.jpg...）
  └── labels/        ← 里面全是标签（001.txt, 002.txt...）

data.yaml

path: /path/to/my_dataset  # 改为 my_dataset 的绝对路径
train: images              # 框架会自动去对应的 labels 文件夹找标签
val: images                # 验证集也指向同一个地方

nc: 1                      # 您的类别数量
# 类别名称列表
names:
  0: pole

训练脚本(数据集->模型)

使用CPU训练,训练只支持CUDA,所以AMD显卡无法训练.

from ultralytics import YOLO

def train_pole_detector():
    """
    训练灯杆检测模型
    """
    # 加载预训练模型
    model = YOLO("yolov8n.pt")

    # 开始训练（CPU）
    results = model.train(
        data="./data.yaml",      # 数据集配置
        epochs=100,
        imgsz=640,
        batch=8,
        device="cpu",            # 使用 CPU 训练
        workers=0,               # Windows 建议 0，避免多进程问题
        patience=20,
        lr0=0.001,
        augment=True,
        name="pole_detect",
    )

    return results


if __name__ == "__main__":
    train_pole_detector()

得到最终模型

跑模型测试推理

detect.py

import torch_directml
from ultralytics import YOLO
import cv2

def detect_lamp_pole(image_path):
    """
    使用 YOLO 检测灯杆
    """
    # 加载模型
    model = YOLO("pole.pt")  # 或 "yolo11n.pt"
    
    print(model.names)
    # GPU模式：使用RX580加速
    dml_device = torch_directml.device()
    model.model.to(dml_device)  # 把底层PyTorch模型迁移到DirectML
    # CPU模式：注销上面两行，改为 dml_device = "cpu"
    
    # 推理（device参数可省略，ultralytics会自动检测model.model所在设备）
    results = model(
        image_path,
        # device=dml_device,  # 可省略：model.model.to()后自动使用GPU
        conf=0.04,           # 置信度阈值
        iou=0.45,            # NMS IoU 阈值
        classes=[0],         # 可选：限制类别，person=0
    )
    
    # 显示结果
    for result in results:
        result.show()
        result.save("output.jpg")
    
    return results

# 测试
if __name__ == "__main__":
    detect_lamp_pole("./img/pole2.png")