YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)

YOLOv5 github:GitHub - ultralytics/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite

先从github下载源码到本地,用pycharm打开工程

采集数据:

本次采集数据采用的方式是录取视频,然后用python截取视频帧当做图片,这是处理代码:

import cv2
import os

video_path = "../test"
save_path = "../data/images"

img_name = 0
for video_file in os.listdir(video_path):
    cap = cv2.VideoCapture(os.path.join(video_path,video_file))

    frame_count = 0
    while cap.isOpened():
        ret, frame = cap.read()

        if (frame_count % 12) == 0 :
            print(frame_count)
            if ret:
                frame = cv2.resize(frame,(416,416))
                cv2.imwrite(os.path.join(save_path, str(img_name)+'.jpg'),frame)
                img_name += 1
            else:
                cap.release()
        frame_count += 1

标注数据:

可以使用labelimg工具、make sense(Make Sense)在线标注,注意数据集需要与yolo对应的格式(label文件夹下是txt文件)

这是之前写过的voc数据集转yolo格式的代码:

from xml.etree import ElementTree as ET
import os

xml_path = r"F:\traffic_signs_yolov5\Annotations"
label_path = r"F:\traffic_signs_yolov5\labels\train"
image_path = r"F:\traffic_signs_yolov5\images\train"

class_name = []

for xml_file in os.listdir(xml_path):
    xml2txt_name = os.path.splitext(xml_file)[0] + '.txt'

    tree = ET.parse(os.path.join(xml_path,xml_file))
    root = tree.getroot()
    sign_object = root.findall('object')
    img_size = root.find("size")
    img_width = int(img_size.find("width").text)
    img_height = int(img_size.find("height").text)

    if sign_object:
        for sign in sign_object:
            name = sign.find('name')
            sign_name = name.text
            if sign_name not in class_name:
                class_name.append(sign_name)

            class_id = class_name.index(sign_name)

            bndbox = sign.find('bndbox')
            xmin = float(bndbox[0].text)
            ymin = float(bndbox[1].text)
            xmax = float(bndbox[2].text)
            ymax = float(bndbox[3].text)

            width = (xmax - xmin) / img_width
            height = (ymax - ymin) / img_height
            x_center = (xmin + (xmax - xmin)/2) / img_width
            y_center = (ymin + (ymax - ymin)/2) / img_height

            with open(os.path.join(label_path,xml2txt_name), 'a') as f:
                f.write("{} {} {} {} {}\n".format(class_id, x_center, y_center, width, height))

    else:
        #with open(os.path.join(label_path,xml2txt_name),'a') as f:
        #    f.write("")
        remove_img_name = os.path.splitext(xml_file)[0] + '.jpg'
        if os.path.exists(os.path.join(image_path,remove_img_name)):
            os.remove(os.path.join(image_path,remove_img_name))
            print("remove {}".format(remove_img_name))


print("class_name:",class_name)

最后一个是辅助代码,因为存在一种情况,当我们需要手动删除数据图片时,还需要去找对应的label文件删除,这里写了一个代码,检查labels文件夹中多出来的txt文件并删除:

"""
删除labels文件夹中多出来的label file,labels文件夹与images文件夹需一一对应
"""
import os

images_path = r"F:\my_traffic_sign_data\my_traffic_sign_data\train\images"
labels_path = r"F:\my_traffic_sign_data\my_traffic_sign_data\train\labels"

imgs_name = []

for img_name in os.listdir(images_path):
    imgs_name.append(os.path.splitext(img_name)[0])

for label_file in os.listdir(labels_path):
    if os.path.splitext(label_file)[0] not in imgs_name:
        os.remove(os.path.join(labels_path, label_file))

print(len(os.listdir(images_path)))
print(len(os.listdir(labels_path)))

准备训练:

准备好数据集后,在工程根目录下新建一个datasets文件夹,将我们的数据集放到该文件夹下,如图所示,注意数据集结构:

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第1张图片

新建yaml文件

接着在yolov5-master/data路径下新建一个yaml文件,根据情况修改类别名称(name)与类别数量(nc)如:

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第2张图片

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第3张图片

修改yolov5s.yaml文件

修改yolov5-master/models路径下的yolov5s.yaml文件,只需修改类别的数量

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第4张图片

训练模型:

在开始训练之前,需要安装环境

可以新建一个虚拟环境

conda create --name yolov5 python=3.7 conda activate yolov5

然后找到工程的根目录下的requirements.txt文件,安装所需库

pip install -r requirements.txt # 注意requirement.txt的路径

进入到根目录的train.py文件,找到parse_opt()函数,修改相应参数,cfg:配置文件,data:数据集描述文件,其他参数按需修改

运行该文件开始训练

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第5张图片

借助colab平台训练:

由于在本地用cpu训练非常慢,这里有一个google的colab云平台(需科学上网),可以免费使用gpu训练,但是有限额,过一会就会恢复

https://colab.research.google.com/

需要上传工程,上传大文件会很慢,若需要上传数据集,可以先上传到kaggle,在自己的kaggle账号中新建datasets,虽然也是很慢,但相对快一点,当上传到kaggle完成后,点击copy API command,复制api下载命令

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第6张图片

然后在自己的个人中心中Account,点击生成API Token,然后下载json文件,并将这个json文件上传到colab中

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第7张图片

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第8张图片

然后在colab中运行以下代码:

!pip install kaggle
import json
token = {"username":"ilmcxx","key":"793c1051a6c90a71e37c111bbf3b7dd5"}
#with open('/content/kaggle.json', 'w') as file:
#  json.dump(token, file)
!mkdir -p ~/.kaggle
!cp /content/kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json
!kaggle config set -n path -v /content

然后到colab中复制api下载命令,在前面加上感叹号,代表运行终端命令,运行即可将数据集下载到colab上

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第9张图片

在colab也是需要先安装好环境再进行训练

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第10张图片

应用模型:

训练完成后,会提示相关文件保存在runs/train文件夹下的exp...中

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第11张图片

yolov5-master/runs/train/exp/weights文件夹下的则保存了刚刚训练的最好的模型和最后的模型

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第12张图片

我们打开根目录下的detect.py文件,找到parse_opt()函数,修改相应参数(使用的模型权重文件weights、待识别的图片source、数据集描述文件data),运行该文件

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第13张图片

运行完成后,结果会保存在yolov5-master/runs/detect/exp...文件夹中

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第14张图片

YOLOv5训练自己的数据集(含数据采集、数据标注、数据预处理、借助colab训练)_第15张图片

你可能感兴趣的:(计算机视觉,YOLO,python,深度学习,计算机视觉)