如何用自己的数据训练YOLOv5

最新推荐文章于 2024-07-13 10:27:26 发布

拉普拉斯之妖

最新推荐文章于 2024-07-13 10:27:26 发布

阅读量395

点赞数

文章标签： YOLO python 深度学习

本文链接：https://blog.csdn.net/weixin_45603120/article/details/129661730

版权

如何训练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 ET

dirpath = 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 names

for 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, 0
        sz = root.find('size')
        width = float(sz[0].text)
        height = float(sz[1].text)
        filename = root.find('filename').text
        for child in root.findall('object'):  # 找到图片中的所有框

            sub = child.find('bndbox')  # 找到框的标注值并进行读取
            label = child.find('name').text
            label_ = dict_info.get(label)
            if label_:
                label_ = label_
            else:
                label_ = 0
            xmin = 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_center
                y_center = (ymin + ymax) / (2 * height)
                y_center = '%.6f' % y_center
                w = (xmax - xmin) / width
                w = '%.6f' % w
                h = (ymax - ymin) / height
                h = '%.6f' % h
            except 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:

Go to the official YOLOv5 GitHub repository: https://github.com/ultralytics/yolov5
Click on the “Releases” tab.
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).