YOLOv5训练自己的数据集并优化模型-[从零开始 万字好文]

        以识别水管和电线这两个对象为例，对如何基于yolov5训练自己的数据集进行记录，并附上在学习过程中参考到的一些好文章。

 

目录

YOLOv5训练自己的数据集整个过程主要包括：

环境安装---数据集采集及标注---模型训练---模型推理---模型优化

一、环境安装

二、数据集采集及标注

2.1 数据集采集

2.2 数据标注

2.2.1. 注册

2.2.2. 基本标注

2.3 下载标注结果

2.4 文件整理

 三、模型训练

3.1创建yaml文件

3.2 命令输入

3.3 云服务器训练

3.3.1.拓展库下载安装 remote-ssh 和 remote development

3.3.2.ctrl+shift+p后输入remote，点击remote-ssh connect to host

3.3.3.创建新账户 Add new SSH host，输入要进入的云服务器地址

3.3.4.进入自动识别出来的ip

3.3.5.输入密码

3.3.6.打开远程文件夹

3.3.7.默认的目录不是该服务器的根目录，需要进入根目录的话选择下拉列表的第一个，其他看自己的需要

 3.3.8.继续选择根目录下的下级目录，我这里没选，默认进入整个目录

3.3.9.再次输入密码，即可成功进入服务器​编辑

3.3.10 nohup日志实时监控并打印输出

四、模型推理

五、模型优化

5.1 YOLOv5预标注

5.2 CVAT对预标注结果进行修正

5.2.1.将yolo预标注结果保存成可被cvat读取的格式

5.2.2.CVAT导入数据

5.3.完善标注

5.4 再次训练并检验获得的pt文件精度是否上升

总结

---

YOLOv5训练自己的数据集整个过程主要包括：

环境安装---数据集采集及标注---模型训练---模型推理---模型优化

---

一、环境安装

        这部分不是本博客的重点，很多热门好文也都有保姆级别的教程，这里不多说，还没安装的小伙伴可以参考以下几篇文章:

半小时搞定Yolov5安装配置及使用（详细过程）

深入浅出Yolo系列之Yolov3&Yolov4&Yolov5&Yolox核心基础知识完整讲解

【Yolov5】1.认真总结6000字Yolov5保姆级教程

二、数据集采集及标注

2.1 数据集采集

        我要训练所用的水管及电线在网上图片很少，所以自己动手，将各类水管电线在不同场景下拍摄了近千张。

        而一些常见素材，如人，猫，狗，汽车等素材，可以通过很多资源网站收集，比如：

1.Machine Learning Datasets | Papers With Code

 公开数据集分门别类的整理好了，可以快速找到需要的数据集，以及使用该数据集的一些论文。

2.Academic Torrents

        academictorrents提供超过127TB 的研究数据。一个分布式系统来共享庞大的数据集——供研究人员使用。一个可扩展、安全且容错的数据存储库，具有极快的下载速度。在其他网站上看到的数据，基本都可以通过这个网站下载。

2.2 数据标注

我用的是cvat标注,cvat使用的前提是安装docker，安装教程如下：

CVAT安装及图片标注使用教程，这里copy一下之前我写的cvat使用教程：

2.2.1. CVAT注册

• 第一次使用CVAT时需要注册，
• 用户注册完就可以去创建标注任务了

2.2.2. 基本标注

2.2.2.1 创建Task

• 登录后会看到如下图界面，CVAT的标注最小单位是Task，每个Task为一个标注任务。点击Task按钮可以看到已创建的task, 点击创建新task。
• 创建新任务的页面如下，按照如下标识的步骤，第一步填写任务名称，要关联的Project名称（可选），第二步设置标签，第三步选择标注数据来源，第四步选择高级设置参数，第五步提交。
• 
• label 标签详解
  • 添加label标签有两种模式：Raw和Constructor，Constructor是可以逐一添加、调整label设置的，设置新增标签名称、颜色、属性等。假如要使用筛选功能，可在此处添加“selected”标签，后续可根据此标签对标注数据进行筛选。

  • 

  • Raw则是可以一键复制所以label设置的json格式的文件，此功能在多个任务对应标签一致时，可以很方便的设置多个标签。Raw示例内容如下图所示，json内容可以通过Copy按钮复制内容用于新任务label设置，粘贴后点击Done保存当前设置，Reset则是维持原来的设置，比如新任务中某标签颜色改变，更改后发现新设置不合适，Reset可以一键还原原始设置。

• 高级配置 Advanced configuration很多参数目前阶段用不到，以下仅对几个常用的参数进行说明：

  • Use zip chunks: 针对视频数据，压缩成zip格式
  • Use cache: 数据缓存
  • Image quality: 用于指定上传的图片质量，当高分辨率图片上传太慢时可以降低分辨率提速，但上一步选择从服务器读取数据的话，图片读取速度会大大提升。
  • Overlap size：把目前的数据分成N个Job进行标注时，每个Job之间重复的帧数；
  • Segment size: 把目前的数据分为N个Job，N就是这里的值，主要用于图片太多，需要多个人一起标注的场景，每个人标一个Job，用此参数将数据分块；
  • Start frame: 采集的视频数据可能并不需要全部标注，比如目标帧可能仅仅只是其中第20-100帧，此参数用于设置从第几帧开始标；
  • Stop frame: 跟Start frame对应，这是结束帧位置；
  • Frame step: 可能不需要每帧都标，比如每隔M帧标一帧，M即是此处的值；
  • 其他参数暂不需要，详见官方文档

2.2.2.2 Task列表页面

• 创建好的Task长这样
  • 
  • Actions中展开：
    • Upload annotations: 支持标签导入，比如在Crowdhuman公开数据集已有的人头和人体标签前提下，想继续标注此数据集中的车辆，用于人头、行人、车三类目标的训练集，就可以使用此功能将人头和人体标签导入，仅用CVAT标注车辆标签即可。具体操作移步本教程第5章“标签导入”。
    • Export task dataset: 下载数据集标签到本地。
    • Automatic annotation: 半自动标注，CVAT支持用预训练模型对数据集在线生成标签，如此可以大大提升标注效率，详见第6章“半自动标注”。
    • Export task: 支持导入或导出一个已有的task, 详见此处。
    • Move to project: 顾名思义，移动本task到另一个project中，移动过程中出现的标签属性变化和标签不匹配问题需要重新对应修改标签设置。
    • Delete: 删除本task。 

2.2.2.3 Task内页

• Task内页如下图所示，在Task内页中，可以对task的一些基本信息进行修改，也可以对任务进行分配，多人协同标注时，查看标注进度等。说明如下：
  • ​
    1. 任务名称，可修改
    2. 此Actions同2.2节的Actions说明
    3. Issue Tracker，一般不用，暂时不管
    4. 标签设置，可修改
    5. 将此Task分配给某个成员
    6. 将此Job分配给某个成员进行标注
    7. 确定质检员，由谁来质检标注数据质量

2.2.2.4 标注页面

• 点击相应Job，就会进入标注界面。界面分为5个部分，如下图所示：
  • ​
  • Header: 主要包括CVAT的几个基本功能导航，以及当前用户相关设置, 基本功能包括Project\Tasks\Cloud Storages\Models，分别表示标注项目、标注任务、云存储、预训练模型。
  • Top panel: 主要包括图片导航（选择图片）、Menu、保存、撤销/恢复、全屏、标注结果统计、mode切换。
  • Workspace: 工作区，图像所在区域。
  • Controls sidebar: 控制相关命令侧边栏，包括设置图片大小、位置，创建shape、编辑tracks等标注基本功能。
  • Object sidebar: 标签相关侧边栏，包括两个选项（objects和labels），还有 appearance 相关设置（就是各类标签的展示颜色、透明度等）。

• Navigation 相关
  • CVAT提供选择上一帧、下一帧、以step为单位，跳转至下（上）一帧，顺序播放每一帧、跳转至开始、结尾帧等功能。
  • 选中帧后，要选择图像的位置以及缩放有以下几个基本操作：
    • 移动图像/选择中心位置：使用第一张图的选项。
    • 另外两个功能是：使图像适应屏幕大小/选择一块区域放大。

    • ​

2.2.2.5 标注小技巧

• 标注过程中的快捷键可以省很多力气，以下是比较常用的几个，更多请见官方文档此处：
  • 保存： Ctrl + S
  • 删除： Delete
  • 缩放：鼠标滚轴，往上放大，向下缩小
  • 拖动：鼠标左键点击并拖动图片，调整位置
  • 下一帧：F
  • 上一帧：D
  • 复制框：Ctrl + C
  • 粘贴框：Ctrl + V
  • 向后跳十张: Shift+Ctrl + V
  • 向前跳十张: Sh得到ift+Ctrl + C
  • 重复上一次标注框参数：N，比如连续标一类头框，那么每次重复标时按N可以直接标下一个头框
  • 对选定对象连续标注接下来所有帧内此对象：Ctrl + B
  • 转换标签：Ctrl+(0..9)

2.3 下载标注结果

​

​

到这里就完成数据标注了。

2.4 文件整理

        从cvat下载下来的文件夹里，图片和txt文件是分开放置的，如下图：

        我们要将images和txt文件分开放置，才能使其满足能被yolo可识别的存放格式，创建一个文件夹，我这里起名为waterpipe_electricwire6000+（因为有），再在这个文件夹里分别创建名为images和labels的文件夹用来存放一会要被分类的图片和txt文件:

---

        对已经标注好的图片，我们可以通过图片增强的方法实现数量扩充，这里我们对图片RGB三个颜色通道进行排列组合，就可以得到原先六倍数量的数据集，同时生成与之对应的txt文件，以下是实现代码： 

import cv2
import os
import re
from glob import glob
import shutil
from numpy import TooHardError

dir = "/home/lyr/下载/complete_callout/02/"#转换后图片保存位置
ray_path = "/home/lyr/下载/complete_callout/obj_train_data/"#原始图片位置

labels_dir = "/home/lyr/下载/complete_callout/03/"#转换后图片label保存位置
ray_labels_path = "/home/lyr/下载/complete_callout/obj_train_data/"#原始图片labels位置

frames = glob(os.path.join(ray_path, '*.jpg'))

for i, frame in enumerate(frames):
    pattern = re.compile(r'([^<>/\\\|:""\*\?]+)\.\w+$')
    data = pattern.findall(frame)
    data_now = data[0]
    Newdir = os.path.join(ray_path, str(data_now) + '.jpg')
    img = cv2.imread(Newdir)
    # print(Newdir)
    B, G, R = cv2.split(img) #openCV切分为BGR三个通道
    img_GRB = cv2.merge([G, R, B])#重定义为GRB顺序
    img_RGB = cv2.merge([R, G, B])#重定义为RGB顺序
    img_BGR = cv2.merge([B, G, R])#重定义为BGR顺序
    img_GBR = cv2.merge([G, B, R])#重定义为GBR顺序
    img_RBG = cv2.merge([R, B, G])#重定义为RBG顺序
    img_BRG = cv2.merge([B, R, G])#重定义为BRG顺序
    
    lablename= Newdir.split("/")[-1]#获取图片的名字（含文件名后缀）
    # print(lablename)
    lablename = lablename.split(".")[0]#获取图片的名字（不含文件名后缀）
    # print(lablename)
    labelpath = os.path.join(ray_labels_path,lablename)#获取该图片labels的完整路径
    all_dir = labelpath+'.txt'
    print(all_dir)
        
    if os.path.exists(all_dir):  #判断该图片在labels路径下是否有对应的label
        # os.rename(labelpath+'.txt',labelpath+'_RGB'+'.txt')#如果有，就在它移动之后重命名操作
        shutil.copyfile(all_dir,labels_dir + lablename + '_RGB.txt') 
        shutil.copyfile(all_dir,labels_dir + lablename + '_BGR.txt') 
        shutil.copyfile(all_dir,labels_dir + lablename + '_GRB.txt') 
        shutil.copyfile(all_dir,labels_dir + lablename + '_GBR.txt') 
        shutil.copyfile(all_dir,labels_dir + lablename + '_RBG.txt') 
        shutil.copyfile(all_dir,labels_dir + lablename + '_BRG.txt') 
        
        
    else:
        print("The file does not exist")
      
    name_RGB = str(data_now) + "_RGB"+".jpg" 
    name_BGR = str(data_now) + "_BGR"+".jpg"
    name_GRB = str(data_now) + "_GRB"+".jpg"
    name_GBR = str(data_now) + "_GBR"+".jpg"
    name_RBG = str(data_now) + "_RBG"+".jpg"
    name_BRG = str(data_now) + "_BRG"+".jpg"
    
    
    cv2.imwrite(dir + name_RGB, img_RGB)
    cv2.imwrite(dir + name_BGR, img_BGR)    
    cv2.imwrite(dir + name_GRB, img_GRB)
    cv2.imwrite(dir + name_GBR, img_GBR)
    cv2.imwrite(dir + name_RBG, img_RBG)
    cv2.imwrite(dir + name_BRG, img_BRG)

  

              运行成功就可以得到扩充六倍数量的图片和对应的txt文件了，如果想得到按阿拉伯数字命名的文件，可以参考这篇博客：python对文件/图片按生成顺序进行重命名（从指定数字开始）

 

 三、模型训练

3.1创建yaml文件

        yolo目录下/data/coco128.yaml是这样的：

# YOLOv5  by Ultralytics, GPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here (7 MB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
nc: 80  # number of classes
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']  # class names


# Download script/URL (optional)
download: https://ultralytics.com/assets/coco128.zip

        参考这个，我对目标检测的电线-水管归为一类被测物体，依照coco128.yaml创建一个waterpipe-electricwire.yaml文件：

# YOLOv5  by Ultralytics, GPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here (7 MB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: /home/xxx/waterpipe_electricwire6000+/  # 图片文件和txt文件所在的文件夹
train: /home/xxx/waterpipe_electricwire6000+/images/  # 图片文件 train images (relative to 'path') 128 images
val: /home/xxx/waterpipe_electricwire6000+/images/  # 图片文件 val images (relative to 'path') 128 images
# test:  # test images (optional)

# Classes
nc: 1  # number of classes
names: ['waterpipe_electricwire']  # class names

        train和val是训练集和验证集图片的路径，我用的是绝对路径，听说也可以是相对路径，labels的路径不用特别给出来，yolov5可以通过图片路径推断出label路径。nc为类别数。names为类别名称，根据自己的数据集进行修改。

3.2 命令输入

cd yolov5-master
python train.py --img 640 --batch 32 --epochs 300 --data ./data/waterpipe-electricwire.yaml --cfg ./models/yolov5s.yaml --weights '' 

img是输入图片大小，网络会自动按参数进行resize，默认640X640；

batch-size是batch数；

epochs是训练轮数，默认300轮；

waterpipe-electricwire.yaml是数据集的yaml文件，我们存放在data文件夹里；

cfg是模型的yam了l文件，一般存放在models文件夹里；

weights是权重文件 .pt 格式，可以输入空格，代表使用随机权重，或者输入权重文件的路径。

顺利的话能看到训练开始了：

 训练结束：

 到这里就完成模型的初步训练了。

3.3 云服务器训练

        如果有云服务器的话，在云服务器端训练会好很多，不用担心训练占本地内存，下面以用vs code连接云服务为例进行说明，在我之前博客中也有记录：VS code连接云服务器--ubuntu18.4

具体方法为：

3.3.1.拓展库下载安装 remote-ssh 和 remote development

3.3.2.ctrl+shift+p后输入remote，点击remote-ssh connect to host

3.3.3.创建新账户 Add new SSH host，输入要进入的云服务器地址

3.3.4.进入自动识别出来的ip

3.3.5.输入密码

3.3.6.打开远程文件夹

3.3.7.默认的目录不是该服务器的根目录，需要进入根目录的话选择下拉列表的第一个，其他看自己的需要

 3.3.8.继续选择根目录下的下级目录，我这里没选，默认进入整个目录

3.3.9.再次输入密码，即可成功进入服务器

问题说明：

        如果出现无法进入的现象，就是哪一步操作失误了，选close remote重试就好了

3.3.10 nohup日志实时监控并打印输出

        而为了防止本地与云端网络连接不稳定或者主机没电关机的情况出现，我们这里采用后台执行nohup的方式保证程序跑起来不间断

指定log输出路径和文件并记录程序的PID

# nohup xxxx xxxx > xxx(ur_log_file) 2>&1 & echo $! > xxx(pid file)
nohup python train.py > train.log 2>&1 & echo $! > run.pid

        nohup会将所有的输出默认写入在你指定的的log_file文件中，我这里的就是train.log；
同时会将当前python的PID记录在run.pid中，方便你中止程序（一般训练中都用会使用多进程，所以实际上会有多个程序，但是nohup只会记录当前启动的那个程序的PID，不过不影响，因为你只要杀掉1个子进程，所有进程都会中止）。

实时监控nohup输出log的内容：

tail -f train.log 

四、模型推理

        训练结束后，结果会说明生成pt文件的所在位置，如果是在本地训练的话，就可以直接用，如果是云服务器训练出来的，就下载下来放到本地yolo的合适目录下，输入指令：

python detect.py  --weights last.pt

        就可以使用我们自己训练出来的权重文件进行目标检测了。这里我用来进行目标检测的是之前没有经过标注的全新照片，这样更有利于直观地看到权重模型的检测效果。

 

 

可以看到整体的解决效果不错，但还有优化空间。 

五、模型优化

这里模型优化采用YOLOv5预标注+CVAT修正labels的方法，

以下方法在我上一篇博客中有记录：YOLOv5与CVAT联合进行预标注并修正labels--[全网独家]

5.1 YOLOv5预标注

yolov5在模型推理阶段，输入命令：

python detect.py --save-txt

即可获得默认设置下的预标注图像 + txt文件。

        该命令中save_txt选项用于生成结果的txt标注文件，会生成与图片相同文件名的txt检测框信息文件，每个txt文件会生成一行或多行的信息，一行信息对应一个label, 信息包括类别序号（第几类物体）、xcenter ycenter w h（label的中心坐标信息），label的长宽信息，均为归一化数值，如下图：

输入命令：

python detect.py --save-txt  --save-crop

即可获得默认设置下的预标注图像 + txt文件 + 被检物体的截取图像。

 也可以在指定其他信息：

 python detect.py

        --data waterpipe-electricwire.yaml        #指定yaml文件

        --weights last.pt        #权重文件

        --source home/images/        #图片读取位置

        --device 0

        --save-txt

        --save-crop

5.2 CVAT对预标注结果进行修正

需要安装cvat，还没安装的可以看我这篇博客，这里不展开说明：

CVAT安装及图片标注使用教程

5.2.1.将yolo预标注结果保存成可被cvat读取的格式

        创建一个自命名文件夹，自命名文件夹再创建一个obj_train_data文件夹、一个data文件、一个names文件、一个txt文件：

obj_train_data文件用来存放没有检测框的原图 + 被检测生成的txt文件     

这里有个问题要注意：有些图片在预训练时没有被检测到目标，所以没有生成对应的txt文件

而obj_train_data文件中的images和txt是要一一对应的

这就要我们提前手动剔除没有txt文件的图片   

 obj.data写入以下信息，classes = 1即我要检测的类别只有一个

 obj.names文件即是定义类别的具体名称，我这里只有一个

 train.txt文件存放所有图片的位置，但要注意，这里的位置信息不是绝对位置，而是

data/+相对路径

 生成这个信息的代码可以看我这篇博客：python输出文件的位置信息

运行代码会自动生成一个含有图片位置信息的train.txt文件，位于被测图片同级目录下，生成结果也会显示在终端：

5.2.2.CVAT导入数据

打开cvat，点击projects, 创建create a new project

项目随意命名，一定要注意的是，

constructor创建的标签名和上述obj.names文件定义类别的名称是一致的

subimt 后选择open project, 点击import dataset上传材料 

 选择yolo 1.1，将包含 obj_train_data、data文件、names文件、txt文件的压缩包上传

5.3.完善标注

更多使用方法看我的这篇博客： CVAT安装及图片标注使用教程_问题多多快快改的博客-CSDN博客

这里只是简单介绍了yolo+cvat的联合使用，更多方法欢迎评论区留言讨论

5.4 再次训练并检验获得的pt文件精度是否上升

        训练方法上文已经说得很清楚了，检测精度这里就不展开说明了，最直观的方法就是对未被标注过的目标检测物体进行多次识别看是否达到预期效果。

总结

        以上就是yolo如何训练自己数据集的方法，有什么不足之处欢迎大家指出，有什么问题也欢迎评论区留言讨论。