本文分享自华为云社区《YoloV5实战:手把手教物体检测——YoloV5》,作者: AI浩 。  

摘要
YOLOV5严格意义上说并不是YOLO的第五个版本,因为它并没有得到YOLO之父Joe Redmon的认可,但是给出的测试数据总体表现还是不错。详细数据如下
v2-ce5ca0b79ff76765834520f8b9d77ea3_720w.jpg
YOLOv5并不是一个单独的模型,而是一个模型家族,包括了YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x、YOLOv5x+TTA,这点有点儿像EfficientDet。由于没有找到V5的论文,我们也只能从代码去学习它。总体上和YOLOV4差不多,可以认为是YOLOV5的加强版。
项目地址:GitHub - ultralytics/yolov5: YOLOv5 ???? in PyTorch > ONNX > CoreML > TFLite

训练  1、下载代码
项目地址:https://github.com/ultralytics/YOLOv5,最近作者又更新了一些代码。
v2-7a8b4c51f4a2d07705f6fafc6887a1a9_720w.jpg

  2、配置环境
  1. matplotlib>=3.2.2
  2. numpy>=1.18.5
  3. opencv-python>=4.1.2
  4. pillow
  5. PyYAML>=5.3
  6. scipy>=1.4.1
  7. tensorboard>=2.2
  8. torch>=1.6.0
  9. torchvision>=0.7.0
  10. tqdm>=4.41.0
3、准备数据集

数据集采用Labelme标注的数据格式,数据集从RSOD数据集中获取了飞机和油桶两类数据集,并将其转为Labelme标注的数据集。
数据集的地址: https://pan.baidu.com/s/1iTUpvA9_cwx1qiH8zbRmDg
提取码:gr6g
或者:LabelmeData.zip_yolov5实战-深度学习文档类资源-CSDN下载
将下载的数据集解压后放到工程的根目录。为下一步生成测试用的数据集做准备。如下图:
v2-4d1c4efbfd1bc6d54cee585a9d992434_720w.jpg

4、生成数据集
YoloV5的数据集和以前版本的数据集并不相同,我们先看一下转换后的数据集。
数据结构如下图:
v2-b3dc5e3769720ce2b24678e9fbc08e3a_720w.jpg
images文件夹存放train和val的图片
labels里面存放train和val的物体数据,里面的每个txt文件和images里面的图片是一一对应的。
txt文件的内容如下:
v2-4d39f110e388e91dff59249cf02a614d_720w.jpg
格式:物体类别 x y w h
坐标是不是真实的坐标,是将坐标除以宽高后的计算出来的,是相对于宽和高的比例。
下面我们编写生成数据集的代码,新建LabelmeToYoloV5.py,然后写入下面的代码。
  1. import os
  2. import numpy as np
  3. import json
  4. from glob import glob
  5. import cv2
  6. from sklearn.model_selection import train_test_split
  7. from os import getcwd
  8. classes = ["aircraft", "oiltank"]
  9. # 1.标签路径
  10. labelme_path = "LabelmeData/"
  11. isUseTest = True  # 是否创建test集
  12. # 3.获取待处理文件
  13. files = glob(labelme_path + "*.json")
  14. files = [i.replace("\", "/").split("/")[-1].split(".json")[0] for i in files]
  15. print(files)
  16. if isUseTest:
  17.     trainval_files, test_files = train_test_split(files, test_size=0.1, random_state=55)
  18. else:
  19.     trainval_files = files
  20. # split
  21. train_files, val_files = train_test_split(trainval_files, test_size=0.1, random_state=55)
  22. def convert(size, box):
  23.     dw = 1. / (size[0])
  24.     dh = 1. / (size[1])
  25.     x = (box[0] + box[1]) / 2.0 - 1
  26.     y = (box[2] + box[3]) / 2.0 - 1
  27.     w = box[1] - box[0]
  28.     h = box[3] - box[2]
  29.     x = x * dw
  30.     w = w * dw
  31.     y = y * dh
  32.     h = h * dh
  33.     return (x, y, w, h)
  34. wd = getcwd()
  35. print(wd)
  36. def ChangeToYolo5(files, txt_Name):
  37.     if not os.path.exists('tmp/'):
  38.         os.makedirs('tmp/')
  39.     list_file = open('tmp/%s.txt' % (txt_Name), 'w')
  40.     for json_file_ in files:
  41.         json_filename = labelme_path + json_file_ + ".json"
  42.         imagePath = labelme_path + json_file_ + ".jpg"
  43.         list_file.write('%s/%s\n' % (wd, imagePath))
  44.         out_file = open('%s/%s.txt' % (labelme_path, json_file_), 'w')
  45.         json_file = json.load(open(json_filename, "r", encoding="utf-8"))
  46.         height, width, channels = cv2.imread(labelme_path + json_file_ + ".jpg").shape
  47.         for multi in json_file["shapes"]:
  48.             points = np.array(multi["points"])
  49.             xmin = min(points[:, 0]) if min(points[:, 0]) > 0 else 0
  50.             xmax = max(points[:, 0]) if max(points[:, 0]) > 0 else 0
  51.             ymin = min(points[:, 1]) if min(points[:, 1]) > 0 else 0
  52.             ymax = max(points[:, 1]) if max(points[:, 1]) > 0 else 0
  53.             label = multi["label"]
  54.             if xmax <= xmin:
  55.                 pass
  56.             elif ymax <= ymin:
  57.                 pass
  58.             else:
  59.                 cls_id = classes.index(label)
  60.                 b = (float(xmin), float(xmax), float(ymin), float(ymax))
  61.                 bb = convert((width, height), b)
  62.                 out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
  63.                 print(json_filename, xmin, ymin, xmax, ymax, cls_id)
  64. ChangeToYolo5(train_files, "train")
  65. ChangeToYolo5(val_files, "val")
  66. ChangeToYolo5(test_files, "test")
这段代码执行完成会在LabelmeData生成每个图片的txt标注数据,同时在tmp文件夹下面生成训练集、验证集和测试集的txt,txt记录的是图片的路径,为下一步生成YoloV5训练和测试用的数据集做准备。在tmp文件夹下面新建MakeData.py文件,生成最终的结果,目录结构如下图:
v2-efb4de6863f47a81c46f01010e4086b6_720w.jpg
打开MakeData.py,写入下面的代码。
  1. import shutil
  2. import os
  3. file_List = ["train", "val", "test"]
  4. for file in file_List:
  5.     if not os.path.exists('../VOC/images/%s' % file):
  6.         os.makedirs('../VOC/images/%s' % file)
  7.     if not os.path.exists('../VOC/labels/%s' % file):
  8.         os.makedirs('../VOC/labels/%s' % file)
  9.     print(os.path.exists('../tmp/%s.txt' % file))
  10.     f = open('../tmp/%s.txt' % file, 'r')
  11.     lines = f.readlines()
  12.     for line in lines:
  13.         print(line)
  14.         line = "/".join(line.split('/')[-5:]).strip()
  15.         shutil.copy(line, "../VOC/images/%s" % file)
  16.         line = line.replace('JPEGImages', 'labels')
  17.         line = line.replace('jpg', 'txt')
  18.         shutil.copy(line, "../VOC/labels/%s/" % file)
执行完成后就可以生成YoloV5训练使用的数据集了。结果如下:
v2-82d9661f3dc841f207d65825d5924ca1_720w.jpg

  5、修改配置参数  
  1. 打开voc.yaml文件,修改里面的配置参数
  2. train: VOC/images/train/  # 训练集图片的路径
  3. val: VOC/images/val/  # 验证集图片的路径
  4. # number of classes
  5. nc: 2 #检测的类别,本次数据集有两个类别所以写2
  6. # class names
  7. names: ["aircraft", "oiltank"]#类别的名称,和转换数据集时的list对应

6、修改train.py的参数   
  1. cfg参数是YoloV5 模型的配置文件,模型的文件存放在models文件夹下面,按照需求填写不同的文件。
  2. weights参数是YoloV5的预训练模型,和cfg对应,例:cfg配置的是yolov5s.yaml,weights就要配置yolov5s.pt
  3. data是配置数据集的配置文件,我们选用的是voc.yaml,所以配置data/voc.yaml
  4. 修改上面三个参数就可以开始训练了,其他的参数根据自己的需求修改。修改后的参数配置如下:
  5. parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
  6. parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml path')
  7. parser.add_argument('--data', type=str, default='data/voc.yaml', help='data.yaml path')
修改完成后,就可以开始训练了。如下图所示:

v2-c00ddb50aa45d0eb8d8b8102e4aaef7e_720w.jpg

7、查看训练结果
在经历了300epoch训练之后,我们会在runs文件夹下面找到训练好的权重文件和训练过程的一些文件。如图:
v2-cabb03a6cc58c41e966e1208233ff4a4_720w.jpg
v2-bdef1216bb43e208235d9827d230d507_720w.jpg
v2-ee49e3fb04b6198523738f1cfa4d0316_720w.jpg
v2-45745404f13cf1fca66ad674cb897d98_720w.jpg
v2-3da0c3d14a8a20f85eaa82d364f8673e_720w.jpg

测试
首先需要在voc.yaml中增加测试集的路径,打开voc.yaml,在val字段后面增加test: tmp/test.txt这行代码,如图:
v2-8ef99c528bc0aedf3f594499fdb5aa5b_720w.jpg
修改test.py中的参数,下面的这几个参数要修改。
  1. parser = argparse.ArgumentParser(prog='test.py')
  2. parser.add_argument('--weights', nargs='+', type=str, default='runs/exp7/weights/best.pt', help='model.pt path(s)')
  3. parser.add_argument('--data', type=str, default='data/voc.yaml', help='*.data path')
  4. parser.add_argument('--batch-size', type=int, default=2, help='size of each image batch')
  5. parser.add_argument('--save-txt', default='True', action='store_true', help='save results to *.txt')
在275行 修改test的方法,增加保存测试结果的路径。这样测试完成后就可以在inference\images查看到测试的图片,在inference\output中查看到保存的测试结果。

如图:
v2-56d0d0f77fa3958cd7056b90bdbf4dae_720w.jpg

下面是运行的结果:
v2-bb04e4395f49b95b67f8d51b5fc246fe_720w.jpg