基于TensorFlow 的垃圾分拣机器人设计

吴蓬勃， 姚美菱， 王 拓， 孙青华

（石家庄邮电职业技术学院电信工程系，石家庄050021）

0 引 言

随着人工智能技术的飞速发展，许多国家已将人工智能作为国家发展战略。为此，国家制定了《新一代人工智能发展规划》［1］。人工智能领域面临着巨大的人才缺口［1］，各高校相继开设了人工智能相关专业［2］。以TensorFlow为代表的深度学习技术、基于深度学习的机器人控制技术逐渐成为人工智能专业的核心课程。受限于深度学习的高复杂度，部分院校的深度学习课程存在着理论与实践脱节的情况，实验很难开展。深度学习的实验平台价格昂贵，实验内容缺少与实际案例的结合，课堂实践内容缺少对课外实践与学科竞赛的支持［1］。

相关学者针对机器视觉和深度学习技术，研制了一些实验平台。文献［3］中使用π超级计算机和NVIDIA Pascal加速卡，构建了校级高性能计算机平台，并部署了TensorFlow等多款深度学习软件框架，面向校内提供深度学习应用服务［3］。文献［4］中使用基于NVIDIA Jetson TX1研制了计算机视觉创新实验平台，使得计算机视觉可以脱离计算机，应用于嵌入式系统领域；并将传统机器视觉算法和深度学习算法在实验平台上进行了验证。该实验平台具有一定的先进性，但是与实际案例的结合尚有欠缺。文献［5］中使用机器视觉，设计了水果雕花机器人创新实验，案例内容丰富，采用的是传统机器视觉技术，与新技术的对接尚待加强。

本实验平台以社区垃圾分拣为应用场景，基于TensorFlow深度学习框架和机械臂控制技术构建垃圾视觉分拣系统。实现了实验内容与实际应用案例的结合，学生在学习人工智能先进技术的同时，也熟悉了我国的垃圾分类制度，促进了学生垃圾分类习惯的养成。同时，本实验平台所采用的NVIDIA Jetson TX2嵌入式机器学习平台和智能机械臂控制系统，也为学生的课外实践和学科竞赛提供了有效的支持。

1 整体设计方案

本实验平台系统框架如图1所示，主要包括：嵌入式AI计算机模块NVIDIA Jetson TX2、STM32控制板、4轴机械臂、传输带、USB高清摄像头以及垃圾桶等设备。NVIDIA Jetson TX2使用垃圾分类模型对垃圾进行视觉分类。STM32控制板负责机械臂和传输带的运动控制，同时采集垃圾桶满、污水收集桶空、满状态并上传到云端。

图1 系统架构图

系统运行流程为：各类垃圾通过垃圾投放口进入系统，传输带将垃圾逐个分离；传输带末端安装有光电传感器，若垃圾从传送带滑落到机械臂托盘，传输带将停止运转；机械臂将垃圾搬运到摄像头下方，进行图像识别分类；NVIDIA Jetson TX2将垃圾分类结果通过串口发送给STM32，STM32控制机械臂将垃圾搬运、倾倒到指定的垃圾桶。

实验平台的核心处理器选型应遵循：先进性、易用性、实用性、易教学和高性价比。目前深度学习嵌入式实验平台的架构主要包括：CPU＋VPU（例如：树莓派＋Movidius VPU）、CPU ＋GPU（例如：NVIDIA Jetson TX2以及RK3399）、CPU ＋TPU（例如：Google Coral Dev Board：NXP IMX8＋ Google Edge TPU）。

由于数据的采集标注和模型的训练都在PC机或者服务器进行，大部分依赖于NVIDIA显卡。相对于其他方案，NVIDIA Jetson TX2嵌入式AI计算单元可以直接运行在计算机上训练好的模型文件，不需要进行格式转换，不会有运算性能的较大损失。可以实现计算机到嵌入式平台的平滑切换，更加适合深度学习的实验教学。

此外，NVIDIA Jetson TX2嵌入式AI计算单元采用了NVIDIA Pascal GPU架构，配有256个NVIDIA CUDA核心；具有双核Denver2 64 bit CPU和4核ARM A57 Complex，提供丰富的标准硬件接口［6］。使得NVIDIA Jetson TX2具备更加先进的性能，可应用于教学、科研和课外学术竞赛。

2 平台硬件设计

实验平台硬件组成，如图2所示。主要包括：TensorFlow深度学习单元、执行单元和控制、感知与通信单元。

TensorFlow深度学习单元由：NVIDIA Jetson TX2开发板、配套的显示与输入设备、USB摄像头、无线路由器和PC机组成。数据的采集和标注、模型的训练均通过PC机完成，训练好的模型通过无线网络传送到NVIDIA Jetson TX2开发板运行测试，测试结果通过串口输出到STM32开发板。

图2 实验平台硬件框图

控制、感知与通信单元包括：STM32控制板、多路光电检测传感器以及NB-IOT通信单元。STM32控制板通过串口接收垃圾识别结果，并进行数据过滤，通过另外一路串口驱动机械臂和传输带运动。STM32控制板还可以通过多路光电开关检测传输带上的垃圾是否进入机械臂托盘、垃圾桶和污水桶空、满等情况，通过NB-IOT实时传输垃圾桶、机械臂和传输带状态到云端的服务器。

执行单元主要包括：4轴机械臂和传输带。机械臂带有步进电机接口，可直接控制传输带启停。用户只需通过机械臂的串行接口，就可控制机械臂和传输带的运动。

3 平台实验设计

Google的TensorFlow深度学习框架已在计算机视觉、自然语言理解和语音识别等方面得到了广泛应用［7-8］。本实验平台，以垃圾分拣为应用场景，以TensorFlow Object Detection API开源代码库为基础，应用MobileNet SSD模型，对垃圾进行视觉分类，通过STM32控制机械臂实现垃圾分拣。

具体实验流程，如图3所示。主要包括：自定义数据集的制作、公共数据集的裁剪、数据的合并与格式转换、模型训练、模型测试和模型部署几个部分。

图3 实验流程图

3.1 垃圾的图像采集、标注与格式转换

TensoFlow的图像识别模型，需要经过大量垃圾图片的训练，才能取得良好的垃圾识别性能。垃圾的图像数据采集至关重要。

数据集的获取途径主要有两种：自己采集并制作数据集（即：自定义数据集）、采用公共数据集。公共数据集已经进行了良好的图片筛选和预处理，训练效果良好，是垃圾图像数据获取的有效途径。公共数据集所识别的物体种类有限，如果所要识别的垃圾在现有数据集中不存在，则需要手动采集制作数据集。

自定义数据集的制作流程包括：图片数据的采集、图片的预处理以及图像标注。垃圾图片的采集可以通过相机拍摄实物获取，也可通过网络爬虫等工具网络获取。为提高准确率，要求在不同背景、不同光线、不同角度、不同垃圾数量等情况下拍摄垃圾图片。通过网络爬虫等工具获取的图片，需要进行人工筛选，删掉一些不合格的图片。图像的标注采用LableImg软件实现，软件输出为Pascal可视对象数据库（Pascal Visual Object Classes，Pascal VOC）格式的XML标注文件。

对于公共数据集，本平台采用了微软的日常场景通用物体（Miscrsoft Common Objects in COntext，MS COCO）数据集，与Pascal VOC相比，MS COCO数据集具有更多的物体种类（COCO为90类，Pascal VOC为20类），每张图片中包含更多的种类和更多的小物体，可有效提高模型识别物体的精度。但是，MS COCO数据集非常庞大，数据集中的物体种类并不是垃圾识别全部需要的，需要进行数据集的裁剪，提取出需要识别的物体种类的图片和标注信息。

为了实现垃圾图像自定义数据集与公共数据集的合并，需要将从MS COCO数据集中提取的图片数据转换为VOC格式的数据，如图4所示。

图4 MS COCO数据集到VOC数据集格式的转换

MS COCO数据集主要包括3个文件夹：图像标注文件夹annotations，训练用图片文件夹train和验证用图片文件夹val。其中，文件夹annotations中存放了caption、instance和person的JSON格式的标注文件，对于物体识别，只需关注instance的标注文件。VOC数据集包括6个文件夹，分别是：Annotations、ImageSets和 JPEGImages、 labels、 SegmentationsClass 和SegmentationObject。对于物体识别，只需要关注前3个文件夹，分别为：XML格式的标注数据文件夹Annoatations、图片训练验证比例划分文件夹ImageSets和图片数据JPEGImages。

数据格式的转换过程：根据垃圾分类规定，从MS COCO数据集中筛选出相应的物体分类列表；根据分类列表从MS COCO数据集annotations文件夹中的instance＿train2017.json 和instance＿val2017.json 中，提取相应分类的标注信息，并生成相应的XML文件；保存到VOC的Annotations文件夹。根据分类列表从MS COCO数据集train和val文件夹中提取图片放入VOC的JPEGImages文件夹。最后，在VOC的ImageSets文件夹，生成图片训练和验证索引文件。

公共数据集提取的数据与自定义数据集合并后，通过TensorFlow的Objection Detection API库程序creat＿pascal＿tf＿record.py［9］，生成统一的TFRecord 格式文件。

3.2 MobileNet SSD 模型

基于深度学习的目标检测方法主要包括：区域候选模型算法（例如：Faster-RCNN［10］）和回归模型算法（例如：单发多盒探测器（Single Shot MultiBox Detector，SSD）［11］）［12］。后者与前者相比，虽然对小目标的检测精度稍低，但具有更快的运行速度，更适合资源受限的嵌入式系统应用。

SSD模型由Liu W等［11］提出，采用单个深度神经网络多盒检测器对图像中的物体进行检测［9，11］。为使得模型更加适用于嵌入式系统，减少运算量，Andrew等［13］基于SSD框架提出了MobileNet模型，将SSD模型的超分辨率测试序列网络替换成MobileNet网络；使用深度可分离卷积替换原来的标准卷积核，减少了计算量；通过引入两个超参数，减少了参数数量和计算量［12］。

MobileNet SSD模型已经在众多的物体识别场合得到了广泛应用，使用此模型进行教学实验，可使教学实验与实际应用对接。

3.3 模型训练与测试

考虑到教学实验与学科竞赛在时间、硬件、性能要求等方面的区别，本实验分为基础实验和课外扩展实验两部分。基础实验采用自定义数据集进行模型训练，每个学生选择一种垃圾进行模型训练。学生可对数据采集、数据标注、模型训练、模型测试与部署等环节有更加系统全面的认识；模型训练所需时间更短，更易于课堂教学。课外扩展实验，需要采用更高性能的GPU硬件，采用公共数据集与自定义数据集融合的数据集，训练的时间将更长，从而到达接近生产实际的垃圾分拣水平。

模型训练与测试的具体流程为：

（1）下载预训练模型ssd＿mobilenet＿v1＿coco＿2017＿11＿17 到TensorFlow Object Detection API库的object＿detection文件夹。

（2）根据垃圾数据集的实际种类，修改object＿detection／data文件夹中的lable map文件mscoco＿label＿map.pbtxt。

（3）修改object＿detection＼samples＼configs 中的配置文件ssd＿mobilenet＿v1＿coco.config。根据实际训练的垃圾种类，修改num＿classes；指定训练和验证的数据集文件路径input＿path（即：训练和验证的TFRecord文件路径）；指定lable map文件存放路径。

（4）根据计算机CPU和内存配置情况，酌情调整模型每次训练的批次大小batch＿size。

（5）当训练的损失率total loss逐渐减小，准确度mAP逐渐增大并接近1.0时，可停止训练过程。通过TensorFlow Object Detection API库的模型文件输出程序export＿inference＿graph.py，将训练好的模型和参数文件转换为可独立运行的PB模型文件［9］。图5为使用有害垃圾——药片板的图像数据集（319张图片），对模型训练9 817次后，模型的识别率，识别率已经达到了80%。图6为使用药片板数据集对模型训练11 026次的total loss图，损失率降低到了10%。

（6）通过TensorFlow Object Detection API库的测试程序object＿detection＿tutorial.py对训练好的PB 模型文件进行测试，考查模型训练的结果。图7为使用药片板数据集对模型训练11 026次后的识别效果图，可看出已能够识别出全部的药片板。

图5 模型训练9 817次的识别率

图6 模型训练11 026次的整体损失率图

图7 模型训练11 026次后的识别效果图

3.4 垃圾识别模型在TX2上的部署

在计算机上训练好的垃圾识别模型文件，需要在NVIDIA Jeston TX2（后面简称TX2）平台上进行部署。TX2平台上需事先安装好NVIDIA的神经网络推理加速框架TensorRT。TensorRT可以将TensorFlow、Caffe、PyTorch和MXnet等框架的深度学习模型进行解析，并针对NVIDIA设备进行优化和部署加速［14-15］。

实验流程：

（1）将训练好的最新PB模型文件、lable map文件和ssd＿mobilenet＿v1＿coco.config 模型配置文件拷贝到TensorRT的data目录下。

（2）为增加对新模型文件的支持，需要修改TensorRT的utils目录下的egohands＿models.py文件，添加对新模型的支持。

（3）打开TensorRT 中的camera＿tf＿trt.py 文件，修改模型名称DEFAULT＿MODEL和lable map文件路径DEFAULT＿LABELMAP。

（4）修改TensorRT的utils目录下的visualization.py文件，实现对识别结果的串口输出。

经过以上实验步骤，即可完成模型部署。图8～11为对药片板的识别结果，可见训练好的模型在不同光线、不同姿态、不同种类、远距离和背景杂乱情况下，都可以实现对药片板的准确识别。

图8 强光环境下的识别结果

图9 暗光、摆放姿态改变情况下的识别结果

图10 暗光、不同种类药片板识别结果

图11 暗光、远距离、背景杂乱情况下识别结果

3.5 垃圾分拣机器人实验平台

图12 为垃圾分拣机器人实验平台实物图，其中包括：垃圾投放入口、垃圾分离传输带、垃圾承载托盘、USB高清摄像头、4轴机械臂、NVIDIA Jeston TX2、STM32控制板以及智能垃圾桶等设备。表1列出了目前实验平台可以识别的、方便教学的垃圾种类。

图12 垃圾分拣机器人实验平台

表1 垃圾分拣机器人可识别的垃圾种类列表

4 结 语

本文以垃圾分拣为实际工作场景，基于TensorFlow深度学习框架，采用MobileNet SSD模型，通过NVIDIA Jetson TX2机器学习平台和智能机械臂控制系统，实现了垃圾分拣机器人实验平台的设计和研制，并进行了相应的垃圾分类实验设计；为人工智能相关专业的教学、科研和课外学术竞赛提供了一套完整的教学实验平台，具有良好的实用价值。本垃圾分拣机器人在2019世界机器人大会工业机器人设计大赛中，获得一等奖，并获得现场专家和观众的一致好评。