基于垃圾指标检测工作的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 数据集存储文件目录

模型配置与训练

YOLOv5我的项目的下载地址：https://github.com/ultralytics

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文件中。为了装置依赖尽可能不出错，集体对该文件提出以下两点修改意见：

  1. 自己fork该官网我的项目时，numpy的版本要求是1.17，但理论装置时会呈现一些问题。因而，集体倡议将其版本要求批改为1.17.3。
  2. 该依赖文件所给出的cocoapi.git下载的地址，在理论运行时下载的速度是十分迟缓的，而且在Windows零碎下还可能出现异常。因而倡议将相应语句用下列语句进行代替：

git+https://github.com/philferrie...

（2）配置文件设置

YOLOv5我的项目次要须要配置如下两个配置文件：

  1. 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']

  1. 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等形式来进步该模型的指标检测性能。

参考文献

我的项目地址

[1]TACO：https://momodel.cn/workspace/5f0e734c95faedbb53ab3b26?type=app

[2]yolov5：https://momodel.cn/workspace/5f0e7c929fda75fe7f4b01e9?type=app

次要文献

[1]将COCO中的.json文件转成.txt文件：https://www.jianshu.com/p/8ddd8f3fdf73

[2]Train Custom Data：https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data