基于垃圾指标检测工作的 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 文件中。为了装置依赖尽可能不出错，集体对该文件提出以下两点修改意见：
- 自己 fork 该官网我的项目时，numpy 的版本要求是 1.17，但理论装置时会呈现一些问题。因而，集体倡议将其版本要求批改为 1.17.3。
- 该依赖文件所给出的 cocoapi.git 下载的地址，在理论运行时下载的速度是十分迟缓的，而且在 Windows 零碎下还可能出现异常。因而倡议将相应语句用下列语句进行代替：
git+https://github.com/philferrie…
（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 等形式来进步该模型的指标检测性能。
参考文献
我的项目地址
[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