使用yolo5进行数据训练

计算机视觉

发布日期: 2022-07-13

更新日期: 2022-07-16

文章字数: 1.2k

阅读时长: 5 分

阅读次数:

1.下载数据集并进行本地训练的标准流程

克隆github项目

git clone https://github.com/theAIGuysCode/OIDv4_ToolKit 
pip install -r requirements.txt

下载数据集

数据集网站为https://storage.googleapis.com/openimages/web/index.html

搜索对应的类别，这里以Lion为例，下载训练集200张：
```
python main.py downloader --classes Lion --type_csv train --limit 200
```
- 注意：若某个类别有两个单词组成，单词间使用下划线代替：
```
python main.py downloader --classes Bell_pepper --type_csv train --limit 400
```
- 注意：若想同时训练多个类别，可通过如下代码将多个类别移动至同一类别：
```
python main.py downloader --classes Aircraft Weapon --type_csv test --limit 100 --multiclasses 1
```

数据预处理

下载后的数据存储在OID/Dataset/train/，每个类别的文件包括对应的image和label，对应的label为txt文件，打开后为：

Lion 48.0 268.23846499999996 972.8 561.4438100000001

的label为txt文件，打开后为：

Lion 48.0 268.23846499999996 972.8 561.4438100000001

但目前的label格式并不能直接拿来放在yolov5中训练，需要转换成yolov5接受的格式

打开classes.txt，将里面的内容改为需要转换的类别：

在本例中，修改为：Lion

运行指令python convert_annotations.py,转换label格式。
代码如下：

import os
import cv2
import numpy as np
from tqdm import tqdm
import argparse
import fileinput

# function that turns XMin, YMin, XMax, YMax coordinates to normalized yolo format
def convert(filename_str, coords):
    os.chdir("..")
    image = cv2.imread(filename_str + ".jpg")
    coords[2] -= coords[0]
    coords[3] -= coords[1]
    x_diff = int(coords[2]/2)
    y_diff = int(coords[3]/2)
    coords[0] = coords[0]+x_diff
    coords[1] = coords[1]+y_diff
    coords[0] /= int(image.shape[1])
    coords[1] /= int(image.shape[0])
    coords[2] /= int(image.shape[1])
    coords[3] /= int(image.shape[0])
    os.chdir("Label")
    return coords

ROOT_DIR = os.getcwd()

# create dict to map class names to numbers for yolo
classes = {}
with open("classes.txt", "r") as myFile:
    for num, line in enumerate(myFile, 0):
        line = line.rstrip("\n")
        classes[line] = num
    myFile.close()
# step into dataset directory
os.chdir(os.path.join("OID", "Dataset"))
DIRS = os.listdir(os.getcwd())

# for all train, validation and test folders
for DIR in DIRS:
    if os.path.isdir(DIR):
        os.chdir(DIR)
        print("Currently in subdirectory:", DIR)
        
        CLASS_DIRS = os.listdir(os.getcwd())
        # for all class folders step into directory to change annotations
        for CLASS_DIR in CLASS_DIRS:
            if os.path.isdir(CLASS_DIR):
                os.chdir(CLASS_DIR)
                print("Converting annotations for class: ", CLASS_DIR)
                
                # Step into Label folder where annotations are generated
                os.chdir("Label")

                for filename in tqdm(os.listdir(os.getcwd())):
                    filename_str = str.split(filename, ".")[0]
                    if filename.endswith(".txt"):
                        annotations = []
                        with open(filename) as f:
                            for line in f:
                                for class_type in classes:
                                    line = line.replace(class_type, str(classes.get(class_type)))
                                labels = line.split()
                                coords = np.asarray([float(labels[1]), float(labels[2]), float(labels[3]), float(labels[4])])
                                coords = convert(filename_str, coords)
                                labels[1], labels[2], labels[3], labels[4] = coords[0], coords[1], coords[2], coords[3]
                                newline = str(labels[0]) + " " + str(labels[1]) + " " + str(labels[2]) + " " + str(labels[3]) + " " + str(labels[4])
                                line = line.replace(line, newline)
                                annotations.append(line)
                            f.close()
                        os.chdir("..")
                        with open(filename, "w") as outfile:
                            for line in annotations:
                                outfile.write(line)
                                outfile.write("\n")
                            outfile.close()
                        os.chdir("Label")
                os.chdir("..")
                os.chdir("..")
        os.chdir("..")

移动标签位置

数据预处理后，txt和jpg文件是放在一个文件夹下的，仍然不是yolov5要求的文件格式，即label，image文件夹分开存放。首先删除Label文件夹中的文件，但不要删除Label文件夹；

运行代码：
```
import os
import glob
import shutil
from pathlib import Path

p = Path(r"E:\opencv_learning\OIDv4_ToolKit-master\OID\Dataset\validation\Lion")  #  需要统计的文件夹路径，这里我们统计的是桌面上的test文件夹，可以随便改名
q = Path(r"E:\opencv_learning\Lion\valid\Label") #  将提取出的所有txt放到test2文件夹下，如果没有需要先创建一个，可以随便改名
FileList = list(p.glob("**/*.txt"))  # 找出这个文件夹下的所有txt格式的文件，可以自行换成flac、mp4、mp3，全都可以！

for file in FileList:       # 这个的意思是，遍历所有找出的txt文件，并复制到另一个文件夹中
    shutil.copy(file,q)
```
此时将Lion文件夹下的txt文件复制到Label文件夹下，将Label文件夹重命名为label，复制出来到与yolov5项目同一目录级别的文件夹下；接着同上处理img文件，复制出来放在与label文件夹同一级别的目录中，最终目录格式：
- 注意：如果使用shutil的move()剪切方法会报错，具体原因据说是shutil的源码bug，目前仍然没有修复，所以只能用这种比较麻烦的方法构成所需目录。
- 注意，要在train同一目录下创建data.yaml文件，内容参考其他data.yaml文件。
- 注意：如果要得到正确的训练结果，还需要下载同类别的validation文件，张数约为训练集的百分之20，用以上方法整理到与train文件夹同一目录下，test文件同理，最终目录结构：
lion_detect是我自己创建放几张网上找的图片用来测试的。
开始训练

打开yolov5 6.1项目，设置训练参数：