【目标检测】基于YOLOv5的目标检测分类（Pytorch）（自定义数据集训练模型）

灰太狼 2023-03-05 03:10 143阅读 0赞

**前言**

采用YOLOv5训练自己的数据集。

项目源码：[https://github.com/ultralytics/yolov5][https_github.com_ultralytics_yolov5]

数据集制作教程：[https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data][https_github.com_ultralytics_yolov5_wiki_Train-Custom-Data]

--------------------

### 目录 ###

*  一、准备环节
 *  二、数据集制作
 *   *  2.1 my\_dataset.yaml
     *  2.2 darknet 格式
     *  2.3 voc转darknet代码
 *  三、训练
 *   *  3.1 模型选择
     *  3.2 修改配置文件
     *  3.3 train.py
 *  四、运行
 *  五、检测目标

# 一、准备环节 #

计算机环境：Win10 + Python3.8 + cuda10.1  
主要依赖：

matplotlib>=3.2.2
    numpy>=1.18.5
    opencv-python>=4.1.2
    Pillow
    PyYAML>=5.3.1
    scipy>=1.4.1
    torch>=1.7.0
    torchvision>=0.8.1
    tqdm>=4.41.0
    tensorboard>=2.4.1
    seaborn>=0.11.0
    pandas
    coremltools>=4.1
    onnx>=1.8.1
    scikit-learn==0.19.2
    thop  # FLOPS computation
    pycocotools>=2.0  # COCO mAP

2021-5-16更新，当前yolov5的最新版本为5.0，一定要用1.7版本以上的pytorch。

# 二、数据集制作 #

数据集制作：[https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data][https_github.com_ultralytics_yolov5_wiki_Train-Custom-Data]

使用[Labelbox][]或[CVAT][]对图片进行标注，博主一直使用的是labelImg，数据集是已经标注好的，使用的是VOC格式，所以需要将其转换为darknet格式。

*可参考博客：*  
*[VOC数据集格式介绍][VOC]  
[图像标注工具labelImg安装教程及使用方法][labelImg]*

## 2.1 my\_dataset.yaml ##

首先仿照`voc.yaml`制作自己的`my_dataset.yaml`:  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70]  
创建`my_dataset.yaml`之后，建立一个名为`dataset`的文件夹，用于存放训练集（稍后制作训练集存放于此文件夹）![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 1]  
`dataset`目录下，建立如下目录：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 2]

<table> 
 <thead> 
  <tr> 
   <th>文件夹</th> 
   <th>功能</th> 
   <th>备注</th> 
  </tr> 
 </thead> 
 <tbody> 
  <tr> 
   <td>data</td> 
   <td>存放原数据集，数据集格式为VOC。</td> 
   <td>可去掉</td> 
  </tr> 
  <tr> 
   <td>tool</td> 
   <td>存放数据集格式转换代码，将data内的数据集转换成yolov5所需的数据集格式（darknet 格式）。</td> 
   <td></td> 
  </tr> 
  <tr> 
   <td>images</td> 
   <td>存放数据集中的图片。train为训练集，val为验证集，test为测试集。</td> 
   <td></td> 
  </tr> 
  <tr> 
   <td>labels</td> 
   <td>存放数据集中的标签</td> 
   <td></td> 
  </tr> 
 </tbody> 
</table>

## 2.2 darknet 格式 ##

每一张图片需要一个`.txt`文本，如果图片中不含有目标，可以不创建`.txt`文本。

`.jpg`文件和`.txt`文件需要放在不同目录下，命名相同，扩展名不同：

> dataset/image/train/0000001.jpg \# image  
> dataset/label/train/0000001.txt \# label

每一个`.txt`文本的具体格式：

*  一行一个目标。
 *  每一行的格式为：class x\_center y\_center width height。
 *  需要标准化xywh，即每一个参数的范围为0-1。
 *  类别通过数字进行表示，起始值为1。

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 3]

## 2.3 voc转darknet代码 ##

YOLOv5同样使用VOC数据集进行模型训练，那么其中一定包含将voc格式转换darknet格式的脚本。

该代码就写在`get_voc.sh`脚本内，将其内容复制出来对其进行修改。

根据数据集的不同，对classes内的类别进行修改。

# -*- coding:utf-8 -*-
    import xml.etree.ElementTree as ET
    import shutil
    
    sets = ['train', 'val', 'test']
    classes = ['ore carrier', 'bulk cargo carrier', 'general cargo ship', 'container ship', 'fishing boat', 'passenger ship']
    
    
    def convert(size, box):
        dw = 1./(size[0])
        dh = 1./(size[1])
        x = (box[0] + box[1])/2.0 - 1
        y = (box[2] + box[3])/2.0 - 1
        w = box[1] - box[0]
        h = box[3] - box[2]
        x = "{:.8f}".format(x*dw)
        w = "{:.8f}".format(w*dw)
        y = "{:.8f}".format(y*dh)
        h = "{:.8f}".format(h*dh)
        return x, y, w, h
    
    
    def convert_annotation(image_set, image_id):
        in_file = open('../data/Annotations/%s.xml' % image_id)
        out_file = open('../labels/%s/%s.txt' % (image_set, image_id), 'w')
        tree = ET.parse(in_file)
        root = tree.getroot()
        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)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    
    
    if __name__ == '__main__':
        for image_set in sets:
            image_ids = open('../data/ImageSets/Main/%s.txt' % image_set).read().strip().split()
    
            for image_id in image_ids:
                print(image_set + ' ' + image_id)
                convert_annotation(image_set, image_id)
                # 复制图片到指定目录
                shutil.copy('../data/JPEGImages/%s.jpg' % image_id, '../images/%s/%s.jpg' % (image_set, image_id))

# 三、训练 #

## 3.1 模型选择 ##

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 4]  
这里优先考虑准确度，选择了YOLOv5x6模型。

根据自己的需求选择下载。[https://github.com/ultralytics/yolov5/releases][https_github.com_ultralytics_yolov5_releases]

## 3.2 修改配置文件 ##

复制**data/voc.yaml**并重命名为**data/may\_dataset.yaml**，并修改内容。

# Default dataset location is next to YOLOv5:
    # /parent_folder
    # /VOC
    # /yolov5
    # train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
    train: /home/work401/ZouQuchen/warship/images/train/
    val: /home/work401/ZouQuchen/warship/images/val/
    
    # number of classes
    nc: 6
    
    # class names
    names: [ "AircraftCarrier", "Destroyer", "FishingBoat", "GaurdShip", "OilTanker", "SupplyShip" ]

## 3.3 train.py ##

修改主函数中预训练权重文件和数据集的默认路径。  
![在这里插入图片描述][2021051619151423.png]

# 四、运行 #

运行`train.py`开始训练

读取一些基本数据：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 5]

读取图片信息：![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 6]  
开始训练

# 五、检测目标 #

这里使用源码中提供的模型进行测试。

不得不说，YOLOv5的检测效果真是太棒了！

**步骤1**：下载模型或将训练好的模型放在`weights`目录下。（这里使用源码中提供的模型）  
![在这里插入图片描述][20201123203319472.png_pic_left]  
**步骤2**： 运行detect.py进行检测。

这是detect.py的main函数，可修改`weights`选项选择不同的模型；修改`source`选择检测文件的路径。

if __name__ == '__main__':
        parser = argparse.ArgumentParser()
        parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
        parser.add_argument('--source', type=str, default='data/test-images', help='source')  # file/folder, 0 for webcam
        parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
        parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
        parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
        parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
        parser.add_argument('--view-img', action='store_true', help='display results')
        parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
        parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
        parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
        parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
        parser.add_argument('--augment', action='store_true', help='augmented inference')
        parser.add_argument('--update', action='store_true', help='update all models')
        parser.add_argument('--project', default='runs/detect', help='save results to project/name')
        parser.add_argument('--name', default='exp', help='save results to project/name')
        parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
        opt = parser.parse_args()
        print(opt)
    
        with torch.no_grad():
            if opt.update:  # update all models (to fix SourceChangeWarning)
                for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
                    detect()
                    strip_optimizer(opt.weights)
            else:
                detect()

--------------------

**4.2 Colaboratory运行（未成功，有待更新）**

将项目文件（包括数据集）复制到Google云盘上。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left]  
启动Colaboratory，设置GPU运行：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left 1]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left 2]  
挂载：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_center]  
cd到项目目录下：

import os
    os.chdir("/content/gdrive/My Drive/Colab Notebooks/yolov5-master/")

安装依赖（只成功安装了PyYaml，可能是其他的依赖这个平台都有）

!pip install -qr requirements.txt

![在这里插入图片描述][20200831162929952.png_pic_center]

查看配置：

import torch
    from IPython.display import Image, clear_output  # to display images
    from utils.google_utils import gdrive_download  # to download models/datasets
    
    clear_output()
    print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left 3]  
运行：

!python train.py

运行失败，但是在笔记本上可以成功运行这一步。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left 4]  
寻找解决方法中……

[https_github.com_ultralytics_yolov5]: https://github.com/ultralytics/yolov5
[https_github.com_ultralytics_yolov5_wiki_Train-Custom-Data]: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
[Labelbox]: https://labelbox.com/
[CVAT]: https://github.com/opencv/cvat
[VOC]: https://blog.csdn.net/See_Star/article/details/103546650
[labelImg]: https://blog.csdn.net/See_Star/article/details/103324456
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70]: /images/20230208/eb5b90627a734e5eb59356cd6a56630c.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 1]: /images/20230208/213acd880b424bbc94e389cd18fe396b.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 2]: /images/20230208/6792dcbb0d474ce2a5e37e5410ffb9fc.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 3]: /images/20230208/ab56a7ab36244cd7b080222fe88a19f2.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 4]: /images/20230208/a5b12e74eeb44fd39dd6925ccac5ebd5.png
[https_github.com_ultralytics_yolov5_releases]: https://github.com/ultralytics/yolov5/releases
[2021051619151423.png]: /images/20230208/f8da79f3313e42a7823a3f9d8b718a58.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 5]: /images/20230208/38aa99120a3a472ba14273fa7b5e7236.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70 6]: /images/20230208/b3225df2037b4a5c95a6151f1269dd4e.png
[20201123203319472.png_pic_left]: /images/20230208/613480e4e2e7483e95f7c0182089a0d0.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left]: /images/20230208/e0fe5127eb8241ef9f01dae184146ff4.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left 1]: /images/20230208/53853113844e4462a5ea01281b4fd20a.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left 2]: /images/20230208/746289a0bba24e09a89b9f5de12047b2.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_center]: /images/20230208/aa5dc8f9a6514e99a5c2b93a5eedce7e.png
[20200831162929952.png_pic_center]: /images/20230208/c895bbdd113048788dcf5514c7fc2509.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left 3]: /images/20230208/b655b6cf48d243218a42b77e8826d325.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1NlZV9TdGFy_size_16_color_FFFFFF_t_70_pic_left 4]: /images/20230208/67ef86bcaca54d7ea98597a3c371d927.png