如何用自己的数据训练YOLOv5

如何训练YOLOv5

1. Clone the YOLOv5 repository and install dependencies:

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

2. 整理数据，使其适配YOLO训练

Step1：Organize your dataset in the following format:

dataset
│
└───train
│   └───images
│   │   │   img1.jpg
│   │   │   img2.jpg
│   │   │   ...
│   └───labels
│       │   img1.txt
│       │   img2.txt
│       │   ...
│
└───valid└───images│   │   img1.jpg│   │   img2.jpg│   │   ...└───labels│   img1.txt│   img2.txt│   ...

Step 2： 将xml格式的label转化为txt格式（适配）

# -*- coding: utf-8 -*-
# 需要修改的地方：1. dirpath 2. newdir 3. dict_info 
import os
import xml.etree.ElementTree as ETdirpath = r'D:\2023\SemiDistill\Data\Annotations'  # 原来存放xml文件的目录
newdir = r'D:\2023\SemiDistill\Data\labels'  # 修改label后形成的txt目录if not os.path.exists(newdir):os.makedirs(newdir)dict_info = {'pocket': 0}  # 有几个 类别 填写几个label namesfor fp in os.listdir(dirpath):if fp.endswith('.xml'):root = ET.parse(os.path.join(dirpath, fp)).getroot()xmin, ymin, xmax, ymax = 0, 0, 0, 0sz = root.find('size')width = float(sz[0].text)height = float(sz[1].text)filename = root.find('filename').textfor child in root.findall('object'):  # 找到图片中的所有框sub = child.find('bndbox')  # 找到框的标注值并进行读取label = child.find('name').textlabel_ = dict_info.get(label)if label_:label_ = label_else:label_ = 0xmin = float(sub[0].text)ymin = float(sub[1].text)xmax = float(sub[2].text)ymax = float(sub[3].text)try:  # 转换成yolov3的标签格式，需要归一化到（0-1）的范围内x_center = (xmin + xmax) / (2 * width)x_center = '%.6f' % x_centery_center = (ymin + ymax) / (2 * height)y_center = '%.6f' % y_centerw = (xmax - xmin) / widthw = '%.6f' % wh = (ymax - ymin) / heighth = '%.6f' % hexcept ZeroDivisionError:print(filename, '的 width有问题')with open(os.path.join(newdir, fp.split('.xml')[0] + '.txt'), 'a+') as f:f.write(' '.join([str(label_), str(x_center), str(y_center), str(w), str(h) + '\n']))
print('ok')

3. Create a YAML file:

Create a YAML file (e.g., my_data.yaml) to describe your dataset:

# 需要修改的地方：train、val、names, nc
train: D:\\2023\\SemiDistill\\Data\\ImageSets\\Main\\train # train文件夹路径
val: D:\\2023\\SemiDistill\\Data\\ImageSets\\Main\\val # val文件夹路径
# number of classes
nc: 1
# class names
names: ["pocket"]

4. Train YOLOv5s

python train.py --img <img_size> --batch <batch_size> --epochs <num_epochs> --data <my_data.yaml> --cfg models/yolov5s.yaml --weights <yolov5s.pt> --name yolov5s_results

在yolov5文件夹下terminal执行以上命令，注意修改<>内数据，其中<my_data.yaml>为自己创建的yaml文件路径，<yolov5s.pt>是下载的yolov5s.pt（预权重）文件路径。
pt文件下载方法：
Here are the steps to download the YOLOv5s.pt file:

1. Go to the official YOLOv5 GitHub repository: https://github.com/ultralytics/yolov5

2. Click on the “Releases” tab.
   

3. In the “Assets” section, you will find the pre-trained weights for YOLOv5s in the form of a .pt file. The filename is “yolov5s.pt”.

5. Evaluate your trained model

After training, you will find the model weights in the runs/train/yolov5s_results/weights folder. To test the trained model on your validation dataset, you can use the test.py script.

python test.py --weights runs/train/yolov5s_results/weights/best.pt --data my_data.yaml --img <img_size> --iou-thres 0.65 --conf-thres 0.001

After running the test, you will find the results in the runs/test folder. The results will include metrics such as precision, recall, and mAP (mean average precision).