YOLOv5从零到一:手把手教你构建与训练专属数据集
1. 环境准备与数据收集第一次接触YOLOv5时最让人头疼的就是如何把一堆原始图片变成模型能吃的数据。这就像要给挑食的孩子准备便当得先了解他的口味。YOLOv5支持的数据格式主要有三种VOC、COCO和YOLO格式我们这次重点讲最常用的YOLO格式。我建议先创建一个专门的项目目录比如yolov5_project里面再建两个子文件夹raw_images存放原始图片建议用JPG或PNG格式raw_labels准备存放标注文件注意图片命名不要用中文或特殊字符最好用连续数字或字母编号比如img_001.jpg实测中发现图片尺寸差异过大会影响训练效果。可以用这个Python脚本批量调整大小from PIL import Image import os input_dir raw_images output_dir resized_images target_size (640, 640) # YOLOv5默认输入尺寸 if not os.path.exists(output_dir): os.makedirs(output_dir) for filename in os.listdir(input_dir): img Image.open(f{input_dir}/{filename}) img img.resize(target_size, Image.BILINEAR) img.save(f{output_dir}/{filename})2. 数据标注实战技巧2.1 标注工具选型推荐三款主流标注工具LabelImg新手友好支持VOC和YOLO格式导出CVAT功能强大适合团队协作Roboflow在线工具带自动标注功能以LabelImg为例安装命令很简单pip install labelImg labelImg # 启动图形界面标注时容易踩的坑标注框要完全包裹目标物体但不要留太多空隙不同类别的物体要用不同标签如cat、dog遇到遮挡物体时按可见部分标注2.2 标注文件格式转换YOLOv5需要的标注格式是每张图片对应一个.txt文件内容类似0 0.5 0.5 0.3 0.4 # 类别ID x_center y_center width height这个转换脚本可以把LabelImg的XML转成YOLO格式import xml.etree.ElementTree as ET import os classes [cat, dog] # 你的类别列表 def convert(size, box): dw 1./size[0] dh 1./size[1] x (box[0] box[1])/2.0 y (box[2] box[3])/2.0 w box[1] - box[0] h box[3] - box[2] x x*dw w w*dw y y*dh h h*dh return (x,y,w,h) for xml_file in os.listdir(Annotations): tree ET.parse(fAnnotations/{xml_file}) root tree.getroot() with open(flabels/{xml_file[:-4]}.txt, w) as f: size root.find(size) w int(size.find(width).text) h int(size.find(height).text) for obj in root.iter(object): cls obj.find(name).text if cls not in classes: continue cls_id classes.index(cls) xmlbox obj.find(bndbox) b (float(xmlbox.find(xmin).text), float(xmlbox.find(xmax).text), float(xmlbox.find(ymin).text), float(xmlbox.find(ymax).text)) bb convert((w,h), b) f.write(f{cls_id} { .join([str(a) for a in bb])}\n)3. 数据集划分与配置3.1 划分训练集/验证集建议按8:2的比例拆分这个脚本可以自动划分import os import random image_files [f for f in os.listdir(images) if f.endswith(.jpg)] random.shuffle(image_files) split int(0.8 * len(image_files)) train image_files[:split] val image_files[split:] with open(train.txt, w) as f: for img in train: f.write(fdata/custom/images/{img}\n) with open(val.txt, w) as f: for img in val: f.write(fdata/custom/images/{img}\n)3.2 创建数据集配置文件在yolov5/data/下新建custom.yaml内容示例train: ../train.txt val: ../val.txt nc: 2 # 类别数量 names: [cat, dog] # 类别名称4. 模型训练与调优4.1 选择预训练模型YOLOv5提供四种规格的模型yolov5s速度最快精度最低yolov5m平衡型yolov5l高精度yolov5x极致精度新手建议从yolov5s开始python train.py --img 640 --batch 16 --epochs 100 --data data/custom.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt4.2 关键参数解析--img 640输入图像尺寸--batch 16根据GPU显存调整显存不足时减小--epochs 100训练轮次建议50-300--weights yolov5s.pt加载预训练权重训练过程中可以用TensorBoard监控tensorboard --logdirruns5. 常见问题解决方案5.1 CUDA内存不足报错CUDA out of memory时减小--batch-size比如从16降到8减小--img-size比如从640降到512使用--adam优化器比默认SGD省内存5.2 标注文件不匹配确保每个图片都有对应的.txt标注文件标注文件中的类别ID从0开始连续编号坐标值在0-1范围内5.3 训练指标不理想可以尝试增加数据量至少每类100张以上调整学习率--lr 0.01启用数据增强--augment True最后提醒训练完成后别忘了测试模型效果python detect.py --weights runs/exp/weights/best.pt --source test_images/