yolov5s 预训练模型_基于垃圾目标检测任务的YOLOv5初探

基于垃圾目标检测任务的YOLOv5初探

作者：余敏君

研究背景

垃圾分类作为一种有效处理垃圾的科学管理方案，在提高资源利用率、缓解垃圾生产压力以及改善生态环境等方面具有重要意义，是我国社会主义现代化和城市化进程中所必须采取的策略，备受世界各国的迫切关注。2019年以来，随着上海市、杭州市等垃圾分类重点城市有关生活垃圾分类的立法、执法和监督等工作的顺利开展，人们对垃圾分类相关话题的关注度日渐提升，个人垃圾分类的意识也有了很大的提高。但与此同时，由于垃圾的种类极其丰富，个人对垃圾归类的模糊程度普遍较高，因此，垃圾分类自动化的实现显得尤为重要。垃圾目标检测作为垃圾分类自动化的一个重要环节，本文将尝试实现该过程。所谓目标检测，简单来讲就是检测图像中的对象是什么以及在哪里的问题，即"目标分类+定位"问题。

YOLO系列模型作为one-stage类目标检测任务的代表算法，凭借其快速、泛化性能好等特性深受研究者喜爱。在前不久，YOLOv5也在GitHub上正式发布，其立即在网上引发了广泛热议。本文将参见官方提供的使用教程，尝试简单利用YOLOv5网络模型在TACO数据集上实现垃圾目标检测任务。

数据集处理

TACO是一个数据正在不断增长的垃圾对象数据集，其以树林、道路和海滩为拍摄背景，目前包含60个类别的垃圾对象，共1500张图像，近5千份标注。该数据集项目参见：https://github.com/pedropro/TACO

一方面，考虑到该数据集文件存放格式和标签格式要符合YOLOv5模型的相关要求；另一方面，也考虑到该数据集中各垃圾类型对象的样本数量极其不均(如图1所示)，因此，本文首先需要对数据集进行必要的操作。对于数据集的处理代码可参见Mo项目[1]，项目中的_readme.ipynb文件详细给出了相关代码使用的具体操作。

图1 TACO各类别对象的样本数量

对于TACO数据集的处理主要可分为以下两个过程：

(1)由于TACO提供的标签为coco类型(所有信息存放在annotations.json文件中)，项目首先需要将相关标签转换为yolo类型。与此同时，考虑到本文仅对相关网络模型进行初探和个人硬件设备的欠缺，该项目仅挑选出满足

要求的垃圾对象进行相关实验。上述操作的相关代码可以参考博客[1]提供的代码进行修改，该代码可以作为模板用于满足自定义化coco格式转yolo格式需求。

经初步统计，符合要求的垃圾类型(共8类)及其原编号如下所示：

'Clear plastic bottle': 5

'Plastic bottle cap': 7

'Drink can': 12

'Other plastic': 29

'Plastic film': 36

'Other plastic wrapper': 39

'Unlabeled litter': 58

'Cigarette': 59

对于coco格式转yolo格式需求，其主要需完成以下三项任务：

· 存储生成的标签和图像分别至两个文件目录中，同时标签和图像的文件命名要求一致

· 将目标对象原标签集合递增顺序映射至{0-7}空间

· 由于原始标签中位置信息为{top_x, top_y, width, height}，项目需要将其转换为{center_x, center_y, width, height}格式，并将其数值进行归一化操作

其核心代码部分如下所示(cocotoyolo.py)：

# 将垃圾类型原编号映射至{0-7}空间

label_transfer = {5: 0, 7: 1, 12: 2, 29: 3,

36: 4, 39: 5, 58: 6, 59: 7}

class_num = {} # 记录各类型样本数量

img_ids = data_source.getImgIds()

# 遍历每张图片，对标签进行转换

for index, img_id in tqdm.tqdm(enumerate(img_ids), desc='change .json file to .txt file'):

img_info = data_source.loadImgs(img_id)[0]

# 将含文件夹的路径修改为文件名

save_name = img_info['file_name'].replace('/', '_')

# 移去文件扩展名

file_name = save_name.split('.')[0]

# 获取单张图像的宽和高

height = img_info['height']

width = img_info['width']

# 转换所得txt文件存储路径

save_path = save_base_path + file_name + '.txt'

is_exist = False # 记录图片是否包含目标垃圾类型对象

with open(save_path, mode='w') as fp:

# 根据图片编号找出垃圾对象的编号集合

annotation_id = data_source.getAnnIds(img_id)

boxes = np.zeros((0, 5))

if len(annotation_id) == 0: # 集合大小为0

fp.write('')

continue

# 获取coco格式的标签

annotations = data_source.loadAnns(annotation_id)

lines = '' # 记录转换后yolo格式的标签

# 遍历对象标签集

for annotation in annotations:

# 获取垃圾对象的标签

label = coco_labels_inverse[annotation['category_id']]

if label in label_transfer.keys():

# 垃圾类型属于目标垃圾类型则进行格式转换

is_exist = True

box = annotation['bbox']

if box[2] < 1 or box[3] < 1:

# 如果原标签中出现无长或宽数据的情况则跳过

continue

# top_x,top_y,width,height==>cen_x,cen_y,width,height

box[0] = round((box[0] + box[2] / 2) / width, 6)

box[1] = round((box[1] + box[3] / 2) / height, 6)

box[2] = round(box[2] / width, 6)

box[3] = round(box[3] / height, 6)

label = label_transfer[label] # 标签映射

if label not in class_num.keys():

class_num[label] = 0

class_num[label] += 1

lines = lines + str(label) # 先存储标签

for i in box: # 再存储位置信息

lines += ' ' + str(i)

lines += '\n' # 换行

fp.writelines(lines)

if is_exist:

# 存在目标类型对象，则拷贝图像至指定目录

shutil.copy('data/{}'.format(img_info['file_name']), os.path.join(save_image_path, save_name))

else:

# 不存在则删除所生成的标签文件

os.remove(save_path)

(2)完成标签集的生成后，项目需要对其进行样本划分。首先，项目需要将样本集按训练集使用全部样本，测试集样本占总样本0.1比例的要求(其样本图片的数量分别为1086和109，训练集使用全部样本是考虑到个人硬件设备较差的原因)进行划分。其次，考虑网络模型对样本存储目录的要求，项目需要将相应生成的图像和标签存储至相应的文件夹下，其文件目录格式如图2所示。其实现代码参见Mo项目中的sample.py文件。

图2 数据集存储文件目录

模型配置与训练

Mo平台上部署的项目参见同名项目[2]，以下的所有内容可以参考官方教程[2]。该Mo项目的具体操作流程请参见根目录下的_readme.ipynb文件

模型配置

该部分主要涉及相关依赖安装和配置文件设置这两方面内容，接下来，本文将对上述内容进行简要说明。

(1)依赖安装

YOLOv5是由PyTorch深度学习框架搭建而成，因此，我们首先需要在Python中安装PyTorch框架，安装教程可以参见官网的相关内容。此处给出PyTorch最新CPU版本的安装命令。

pip install torch==1.5.1+cpu torchvision==0.6.1+cpu -f https://download.pytorch.org/..._stable.html

除此之外，YOLOv5模型运行还需要安装额外的第三方依赖包。官方已经将其全部放置在requirements.txt文件中。为了安装依赖尽可能不出错，个人对该文件提出以下两点修改意见：

本人fork该官方项目时，numpy的版本要求是1.17，但实际安装时会出现一些问题。因此，个人建议将其版本要求修改为1.17.3。

该依赖文件所给出的cocoapi.git下载的地址，在实际运行时下载的速度是非常缓慢的，而且在Windows系统下还可能出现异常。因此建议将相应语句用下列语句进行代替：

(2)配置文件设置

YOLOv5项目主要需要配置如下两个配置文件：

data路径下的taco.yaml配置文件(自己创建)。该文件主要配置训练集和测试集存储图片的路径以及对象类型信息。对于标签存储路径的配置，项目会自动根据图片存储路径来确定，这也是为什么我们的数据集文件目录需要符合图2所示要求的原因。该配置文件的内容如下所示。

# train and val datasets (image directory or *.txt file with image paths)

train: taco/images/train/

val: taco/images/test/

# number of classes

nc: 8

# class names

names: ['Clear plastic bottle', 'Plastic bottle cap',

'Drink can',

'Other plastic',

'Plastic film', 'Other plastic wrapper',

'Unlabeled litter', 'Cigarette']

models路径下与模型配置相关的配置文件编写。根据官方教程[2]，研究者可以直接修改models目录下已存在的 yolov5_.yaml系列文件来进行相关配置。这里，本人修改的是yolov5s.yaml文件(因为该实验使用yolov5s.pt作为模型预训练权重)。个人只要将该文件中nc属性的值修改为与上述taco.yaml的nc属性值一致即可。

模型训练

一方面，由于个人电脑内存相对较小，直接使用官方提供的训练参数设置将会导致内存爆炸而无法工作。因此，该实验不得不通过减少输入图像规模和batch_size值来促使模型进行训练。另一方面，由于个人电脑GPU性能较差，该实验选择直接使用CPU来训练相关模型，因此模型训练速度相对较慢。考虑到上述两方面的限制，本实验模型训练时所涉及的相关参数配置如表1所示：

表1 模型训练相关的参数配置

命令行参数

参数含义

设置值

--img

统一输入图像规模

320

--batch

每次网络训练输入图像的数量

4

--epochs

整个数据集参与训练的次数

100

--data

数据集配置文件路径

./data/taco.yaml

--cfg

模型配置文件路径

./models/yolov5s.yaml

--device

训练的设备(CPU or GPU)

cpu

--weights

预训练模型的权重文件

yolov5s.pt

模型训练的调用命令如下所示：

python train.py --img 320 --batch 4 --epochs 100 --data ./data/taco.yaml --cfg ./models/yolov5s.yaml --device cpu --weights yolov5s.pt

关于预训练权重的下载，这里就不进行详细介绍，百度一下应该可以找到许多国内下载的资源。在这里，本人在项目根目录下放置了yolov5s的权重文件(yolov5s.pt)以方便研究者训练模型。当然，我们完全可以不使用预训练权重来直接进行模型训练，只要移去上述命令的--weights yolov5s.pt部分即可。

效果展示

根据官方提供的教程，模型训练所生成的各类结果将被自动放置在根目录的runs文件夹下。其中，weights文件夹下将会存放模型训练所生成的效果最好和时间最近的权重文件，我们可以用这些文件完成模型的调用任务；results.txt文件存放着模型训练过程中的各项指标输出，YOLO项目还自动对该输出结果进行了可视化操作，生成了对应的图表图像。本实验各项指标的输出可视化图像如图3所示，上方五张为训练集对应结果，下方五张为测试集对应结果。

图3 模型训练过程各项输出指标的可视化图像

利用detect.py文件我们可以利用生成的模型来检测目标图像中是否存在需要检测的对象。为了方便起见，我们将生成的模型权重best.pt放置在项目根目录下，将需要检测的图像inference/images文件夹下。参考官方文档，我们只需运行如下代码即可展示模型目标检测的效果：

python detect.py --weights best.pt --img 320 --conf 0.4

目标图像检测的效果如图4所示，检测生成的图像位于inference/output文件夹中。其中，batch_1_000048图像中错误检测出了Drink can对象(此处主要与conf设置值的大小有关，个人可以对其进行适当的调整)。

图4 batch_1_000029(左)和batch_1_000048(右)的检测效果

实验小结

本文主要依托垃圾对象目标检测任务，利用新颖的YOLOv5模型对TACO垃圾对象数据集进行了实验。由于个人硬件设备的不足和时间的限制，参考官方给出的一些输出指标效果图和自行实验所得的图像检测结果，可以看到本文最终所得到的模型的性能其实相对来说不算好(YOLO系列模型本身就是需要跑非常久的时间)。感兴趣地朋友可以尝试增大统一的图像规模、增大batch_size和epoch等方式来提高该模型的目标检测性能。

参考文献

项目地址

主要文献