基于TensorFlow 的垃圾分揀機(jī)器人設(shè)計(jì)

吳蓬勃， 姚美菱， 王 拓， 孫青華

（石家莊郵電職業(yè)技術(shù)學(xué)院電信工程系，石家莊050021）

0 引 言

隨著人工智能技術(shù)的飛速發(fā)展，許多國(guó)家已將人工智能作為國(guó)家發(fā)展戰(zhàn)略。為此，國(guó)家制定了《新一代人工智能發(fā)展規(guī)劃》［1］。人工智能領(lǐng)域面臨著巨大的人才缺口［1］，各高校相繼開(kāi)設(shè)了人工智能相關(guān)專業(yè)［2］。以TensorFlow為代表的深度學(xué)習(xí)技術(shù)、基于深度學(xué)習(xí)的機(jī)器人控制技術(shù)逐漸成為人工智能專業(yè)的核心課程。受限于深度學(xué)習(xí)的高復(fù)雜度，部分院校的深度學(xué)習(xí)課程存在著理論與實(shí)踐脫節(jié)的情況，實(shí)驗(yàn)很難開(kāi)展。深度學(xué)習(xí)的實(shí)驗(yàn)平臺(tái)價(jià)格昂貴，實(shí)驗(yàn)內(nèi)容缺少與實(shí)際案例的結(jié)合，課堂實(shí)踐內(nèi)容缺少對(duì)課外實(shí)踐與學(xué)科競(jìng)賽的支持［1］。

相關(guān)學(xué)者針對(duì)機(jī)器視覺(jué)和深度學(xué)習(xí)技術(shù)，研制了一些實(shí)驗(yàn)平臺(tái)。文獻(xiàn)［3］中使用π超級(jí)計(jì)算機(jī)和NVIDIA Pascal加速卡，構(gòu)建了校級(jí)高性能計(jì)算機(jī)平臺(tái)，并部署了TensorFlow等多款深度學(xué)習(xí)軟件框架，面向校內(nèi)提供深度學(xué)習(xí)應(yīng)用服務(wù)［3］。文獻(xiàn)［4］中使用基于NVIDIA Jetson TX1研制了計(jì)算機(jī)視覺(jué)創(chuàng)新實(shí)驗(yàn)平臺(tái)，使得計(jì)算機(jī)視覺(jué)可以脫離計(jì)算機(jī)，應(yīng)用于嵌入式系統(tǒng)領(lǐng)域；并將傳統(tǒng)機(jī)器視覺(jué)算法和深度學(xué)習(xí)算法在實(shí)驗(yàn)平臺(tái)上進(jìn)行了驗(yàn)證。該實(shí)驗(yàn)平臺(tái)具有一定的先進(jìn)性，但是與實(shí)際案例的結(jié)合尚有欠缺。文獻(xiàn)［5］中使用機(jī)器視覺(jué)，設(shè)計(jì)了水果雕花機(jī)器人創(chuàng)新實(shí)驗(yàn)，案例內(nèi)容豐富，采用的是傳統(tǒng)機(jī)器視覺(jué)技術(shù)，與新技術(shù)的對(duì)接尚待加強(qiáng)。

本實(shí)驗(yàn)平臺(tái)以社區(qū)垃圾分揀為應(yīng)用場(chǎng)景，基于TensorFlow深度學(xué)習(xí)框架和機(jī)械臂控制技術(shù)構(gòu)建垃圾視覺(jué)分揀系統(tǒng)。實(shí)現(xiàn)了實(shí)驗(yàn)內(nèi)容與實(shí)際應(yīng)用案例的結(jié)合，學(xué)生在學(xué)習(xí)人工智能先進(jìn)技術(shù)的同時(shí)，也熟悉了我國(guó)的垃圾分類制度，促進(jìn)了學(xué)生垃圾分類習(xí)慣的養(yǎng)成。同時(shí)，本實(shí)驗(yàn)平臺(tái)所采用的NVIDIA Jetson TX2嵌入式機(jī)器學(xué)習(xí)平臺(tái)和智能機(jī)械臂控制系統(tǒng)，也為學(xué)生的課外實(shí)踐和學(xué)科競(jìng)賽提供了有效的支持。

1 整體設(shè)計(jì)方案

本實(shí)驗(yàn)平臺(tái)系統(tǒng)框架如圖1所示，主要包括：嵌入式AI計(jì)算機(jī)模塊NVIDIA Jetson TX2、STM32控制板、4軸機(jī)械臂、傳輸帶、USB高清攝像頭以及垃圾桶等設(shè)備。NVIDIA Jetson TX2使用垃圾分類模型對(duì)垃圾進(jìn)行視覺(jué)分類。STM32控制板負(fù)責(zé)機(jī)械臂和傳輸帶的運(yùn)動(dòng)控制，同時(shí)采集垃圾桶滿、污水收集桶空、滿狀態(tài)并上傳到云端。

圖1 系統(tǒng)架構(gòu)圖

系統(tǒng)運(yùn)行流程為：各類垃圾通過(guò)垃圾投放口進(jìn)入系統(tǒng)，傳輸帶將垃圾逐個(gè)分離；傳輸帶末端安裝有光電傳感器，若垃圾從傳送帶滑落到機(jī)械臂托盤(pán)，傳輸帶將停止運(yùn)轉(zhuǎn)；機(jī)械臂將垃圾搬運(yùn)到攝像頭下方，進(jìn)行圖像識(shí)別分類；NVIDIA Jetson TX2將垃圾分類結(jié)果通過(guò)串口發(fā)送給STM32，STM32控制機(jī)械臂將垃圾搬運(yùn)、傾倒到指定的垃圾桶。

實(shí)驗(yàn)平臺(tái)的核心處理器選型應(yīng)遵循：先進(jìn)性、易用性、實(shí)用性、易教學(xué)和高性價(jià)比。目前深度學(xué)習(xí)嵌入式實(shí)驗(yàn)平臺(tái)的架構(gòu)主要包括：CPU＋VPU（例如：樹(shù)莓派＋Movidius VPU）、CPU ＋GPU（例如：NVIDIA Jetson TX2以及RK3399）、CPU ＋TPU（例如：Google Coral Dev Board：NXP IMX8＋ Google Edge TPU）。

由于數(shù)據(jù)的采集標(biāo)注和模型的訓(xùn)練都在PC機(jī)或者服務(wù)器進(jìn)行，大部分依賴于NVIDIA顯卡。相對(duì)于其他方案，NVIDIA Jetson TX2嵌入式AI計(jì)算單元可以直接運(yùn)行在計(jì)算機(jī)上訓(xùn)練好的模型文件，不需要進(jìn)行格式轉(zhuǎn)換，不會(huì)有運(yùn)算性能的較大損失。可以實(shí)現(xiàn)計(jì)算機(jī)到嵌入式平臺(tái)的平滑切換，更加適合深度學(xué)習(xí)的實(shí)驗(yàn)教學(xué)。

此外，NVIDIA Jetson TX2嵌入式AI計(jì)算單元采用了NVIDIA Pascal GPU架構(gòu)，配有256個(gè)NVIDIA CUDA核心；具有雙核Denver2 64 bit CPU和4核ARM A57 Complex，提供豐富的標(biāo)準(zhǔn)硬件接口［6］。使得NVIDIA Jetson TX2具備更加先進(jìn)的性能，可應(yīng)用于教學(xué)、科研和課外學(xué)術(shù)競(jìng)賽。

2 平臺(tái)硬件設(shè)計(jì)

實(shí)驗(yàn)平臺(tái)硬件組成，如圖2所示。主要包括：TensorFlow深度學(xué)習(xí)單元、執(zhí)行單元和控制、感知與通信單元。

TensorFlow深度學(xué)習(xí)單元由：NVIDIA Jetson TX2開(kāi)發(fā)板、配套的顯示與輸入設(shè)備、USB攝像頭、無(wú)線路由器和PC機(jī)組成。數(shù)據(jù)的采集和標(biāo)注、模型的訓(xùn)練均通過(guò)PC機(jī)完成，訓(xùn)練好的模型通過(guò)無(wú)線網(wǎng)絡(luò)傳送到NVIDIA Jetson TX2開(kāi)發(fā)板運(yùn)行測(cè)試，測(cè)試結(jié)果通過(guò)串口輸出到STM32開(kāi)發(fā)板。

圖2 實(shí)驗(yàn)平臺(tái)硬件框圖

控制、感知與通信單元包括：STM32控制板、多路光電檢測(cè)傳感器以及NB-IOT通信單元。STM32控制板通過(guò)串口接收垃圾識(shí)別結(jié)果，并進(jìn)行數(shù)據(jù)過(guò)濾，通過(guò)另外一路串口驅(qū)動(dòng)機(jī)械臂和傳輸帶運(yùn)動(dòng)。STM32控制板還可以通過(guò)多路光電開(kāi)關(guān)檢測(cè)傳輸帶上的垃圾是否進(jìn)入機(jī)械臂托盤(pán)、垃圾桶和污水桶空、滿等情況，通過(guò)NB-IOT實(shí)時(shí)傳輸垃圾桶、機(jī)械臂和傳輸帶狀態(tài)到云端的服務(wù)器。

執(zhí)行單元主要包括：4軸機(jī)械臂和傳輸帶。機(jī)械臂帶有步進(jìn)電機(jī)接口，可直接控制傳輸帶啟停。用戶只需通過(guò)機(jī)械臂的串行接口，就可控制機(jī)械臂和傳輸帶的運(yùn)動(dòng)。

3 平臺(tái)實(shí)驗(yàn)設(shè)計(jì)

Google的TensorFlow深度學(xué)習(xí)框架已在計(jì)算機(jī)視覺(jué)、自然語(yǔ)言理解和語(yǔ)音識(shí)別等方面得到了廣泛應(yīng)用［7-8］。本實(shí)驗(yàn)平臺(tái)，以垃圾分揀為應(yīng)用場(chǎng)景，以TensorFlow Object Detection API開(kāi)源代碼庫(kù)為基礎(chǔ)，應(yīng)用MobileNet SSD模型，對(duì)垃圾進(jìn)行視覺(jué)分類，通過(guò)STM32控制機(jī)械臂實(shí)現(xiàn)垃圾分揀。

具體實(shí)驗(yàn)流程，如圖3所示。主要包括：自定義數(shù)據(jù)集的制作、公共數(shù)據(jù)集的裁剪、數(shù)據(jù)的合并與格式轉(zhuǎn)換、模型訓(xùn)練、模型測(cè)試和模型部署幾個(gè)部分。

圖3 實(shí)驗(yàn)流程圖

3.1 垃圾的圖像采集、標(biāo)注與格式轉(zhuǎn)換

TensoFlow的圖像識(shí)別模型，需要經(jīng)過(guò)大量垃圾圖片的訓(xùn)練，才能取得良好的垃圾識(shí)別性能。垃圾的圖像數(shù)據(jù)采集至關(guān)重要。

數(shù)據(jù)集的獲取途徑主要有兩種：自己采集并制作數(shù)據(jù)集（即：自定義數(shù)據(jù)集）、采用公共數(shù)據(jù)集。公共數(shù)據(jù)集已經(jīng)進(jìn)行了良好的圖片篩選和預(yù)處理，訓(xùn)練效果良好，是垃圾圖像數(shù)據(jù)獲取的有效途徑。公共數(shù)據(jù)集所識(shí)別的物體種類有限，如果所要識(shí)別的垃圾在現(xiàn)有數(shù)據(jù)集中不存在，則需要手動(dòng)采集制作數(shù)據(jù)集。

自定義數(shù)據(jù)集的制作流程包括：圖片數(shù)據(jù)的采集、圖片的預(yù)處理以及圖像標(biāo)注。垃圾圖片的采集可以通過(guò)相機(jī)拍攝實(shí)物獲取，也可通過(guò)網(wǎng)絡(luò)爬蟲(chóng)等工具網(wǎng)絡(luò)獲取。為提高準(zhǔn)確率，要求在不同背景、不同光線、不同角度、不同垃圾數(shù)量等情況下拍攝垃圾圖片。通過(guò)網(wǎng)絡(luò)爬蟲(chóng)等工具獲取的圖片，需要進(jìn)行人工篩選，刪掉一些不合格的圖片。圖像的標(biāo)注采用LableImg軟件實(shí)現(xiàn)，軟件輸出為Pascal可視對(duì)象數(shù)據(jù)庫(kù)（Pascal Visual Object Classes，Pascal VOC）格式的XML標(biāo)注文件。

對(duì)于公共數(shù)據(jù)集，本平臺(tái)采用了微軟的日常場(chǎng)景通用物體（Miscrsoft Common Objects in COntext，MS COCO）數(shù)據(jù)集，與Pascal VOC相比，MS COCO數(shù)據(jù)集具有更多的物體種類（COCO為90類，Pascal VOC為20類），每張圖片中包含更多的種類和更多的小物體，可有效提高模型識(shí)別物體的精度。但是，MS COCO數(shù)據(jù)集非常龐大，數(shù)據(jù)集中的物體種類并不是垃圾識(shí)別全部需要的，需要進(jìn)行數(shù)據(jù)集的裁剪，提取出需要識(shí)別的物體種類的圖片和標(biāo)注信息。

為了實(shí)現(xiàn)垃圾圖像自定義數(shù)據(jù)集與公共數(shù)據(jù)集的合并，需要將從MS COCO數(shù)據(jù)集中提取的圖片數(shù)據(jù)轉(zhuǎn)換為VOC格式的數(shù)據(jù)，如圖4所示。

圖4 MS COCO數(shù)據(jù)集到VOC數(shù)據(jù)集格式的轉(zhuǎn)換

MS COCO數(shù)據(jù)集主要包括3個(gè)文件夾：圖像標(biāo)注文件夾annotations，訓(xùn)練用圖片文件夾train和驗(yàn)證用圖片文件夾val。其中，文件夾annotations中存放了caption、instance和person的JSON格式的標(biāo)注文件，對(duì)于物體識(shí)別，只需關(guān)注instance的標(biāo)注文件。VOC數(shù)據(jù)集包括6個(gè)文件夾，分別是：Annotations、ImageSets和 JPEGImages、 labels、 SegmentationsClass 和SegmentationObject。對(duì)于物體識(shí)別，只需要關(guān)注前3個(gè)文件夾，分別為：XML格式的標(biāo)注數(shù)據(jù)文件夾Annoatations、圖片訓(xùn)練驗(yàn)證比例劃分文件夾ImageSets和圖片數(shù)據(jù)JPEGImages。

數(shù)據(jù)格式的轉(zhuǎn)換過(guò)程：根據(jù)垃圾分類規(guī)定，從MS COCO數(shù)據(jù)集中篩選出相應(yīng)的物體分類列表；根據(jù)分類列表從MS COCO數(shù)據(jù)集annotations文件夾中的instance＿train2017.json 和instance＿val2017.json 中，提取相應(yīng)分類的標(biāo)注信息，并生成相應(yīng)的XML文件；保存到VOC的Annotations文件夾。根據(jù)分類列表從MS COCO數(shù)據(jù)集train和val文件夾中提取圖片放入VOC的JPEGImages文件夾。最后，在VOC的ImageSets文件夾，生成圖片訓(xùn)練和驗(yàn)證索引文件。

公共數(shù)據(jù)集提取的數(shù)據(jù)與自定義數(shù)據(jù)集合并后，通過(guò)TensorFlow的Objection Detection API庫(kù)程序creat＿pascal＿tf＿record.py［9］，生成統(tǒng)一的TFRecord 格式文件。

3.2 MobileNet SSD 模型

基于深度學(xué)習(xí)的目標(biāo)檢測(cè)方法主要包括：區(qū)域候選模型算法（例如：Faster-RCNN［10］）和回歸模型算法（例如：?jiǎn)伟l(fā)多盒探測(cè)器（Single Shot MultiBox Detector，SSD）［11］）［12］。后者與前者相比，雖然對(duì)小目標(biāo)的檢測(cè)精度稍低，但具有更快的運(yùn)行速度，更適合資源受限的嵌入式系統(tǒng)應(yīng)用。

SSD模型由Liu W等［11］提出，采用單個(gè)深度神經(jīng)網(wǎng)絡(luò)多盒檢測(cè)器對(duì)圖像中的物體進(jìn)行檢測(cè)［9，11］。為使得模型更加適用于嵌入式系統(tǒng)，減少運(yùn)算量，Andrew等［13］基于SSD框架提出了MobileNet模型，將SSD模型的超分辨率測(cè)試序列網(wǎng)絡(luò)替換成MobileNet網(wǎng)絡(luò)；使用深度可分離卷積替換原來(lái)的標(biāo)準(zhǔn)卷積核，減少了計(jì)算量；通過(guò)引入兩個(gè)超參數(shù)，減少了參數(shù)數(shù)量和計(jì)算量［12］。

MobileNet SSD模型已經(jīng)在眾多的物體識(shí)別場(chǎng)合得到了廣泛應(yīng)用，使用此模型進(jìn)行教學(xué)實(shí)驗(yàn)，可使教學(xué)實(shí)驗(yàn)與實(shí)際應(yīng)用對(duì)接。

3.3 模型訓(xùn)練與測(cè)試

考慮到教學(xué)實(shí)驗(yàn)與學(xué)科競(jìng)賽在時(shí)間、硬件、性能要求等方面的區(qū)別，本實(shí)驗(yàn)分為基礎(chǔ)實(shí)驗(yàn)和課外擴(kuò)展實(shí)驗(yàn)兩部分?；A(chǔ)實(shí)驗(yàn)采用自定義數(shù)據(jù)集進(jìn)行模型訓(xùn)練，每個(gè)學(xué)生選擇一種垃圾進(jìn)行模型訓(xùn)練。學(xué)生可對(duì)數(shù)據(jù)采集、數(shù)據(jù)標(biāo)注、模型訓(xùn)練、模型測(cè)試與部署等環(huán)節(jié)有更加系統(tǒng)全面的認(rèn)識(shí)；模型訓(xùn)練所需時(shí)間更短，更易于課堂教學(xué)。課外擴(kuò)展實(shí)驗(yàn)，需要采用更高性能的GPU硬件，采用公共數(shù)據(jù)集與自定義數(shù)據(jù)集融合的數(shù)據(jù)集，訓(xùn)練的時(shí)間將更長(zhǎng)，從而到達(dá)接近生產(chǎn)實(shí)際的垃圾分揀水平。

模型訓(xùn)練與測(cè)試的具體流程為：

（1）下載預(yù)訓(xùn)練模型ssd＿mobilenet＿v1＿coco＿2017＿11＿17 到TensorFlow Object Detection API庫(kù)的object＿detection文件夾。

（2）根據(jù)垃圾數(shù)據(jù)集的實(shí)際種類，修改object＿detection／data文件夾中的lable map文件mscoco＿label＿map.pbtxt。

（3）修改object＿detection＼samples＼configs 中的配置文件ssd＿mobilenet＿v1＿coco.config。根據(jù)實(shí)際訓(xùn)練的垃圾種類，修改num＿classes；指定訓(xùn)練和驗(yàn)證的數(shù)據(jù)集文件路徑input＿path（即：訓(xùn)練和驗(yàn)證的TFRecord文件路徑）；指定lable map文件存放路徑。

（4）根據(jù)計(jì)算機(jī)CPU和內(nèi)存配置情況，酌情調(diào)整模型每次訓(xùn)練的批次大小batch＿size。

（5）當(dāng)訓(xùn)練的損失率total loss逐漸減小，準(zhǔn)確度mAP逐漸增大并接近1.0時(shí)，可停止訓(xùn)練過(guò)程。通過(guò)TensorFlow Object Detection API庫(kù)的模型文件輸出程序export＿inference＿graph.py，將訓(xùn)練好的模型和參數(shù)文件轉(zhuǎn)換為可獨(dú)立運(yùn)行的PB模型文件［9］。圖5為使用有害垃圾——藥片板的圖像數(shù)據(jù)集（319張圖片），對(duì)模型訓(xùn)練9 817次后，模型的識(shí)別率，識(shí)別率已經(jīng)達(dá)到了80%。圖6為使用藥片板數(shù)據(jù)集對(duì)模型訓(xùn)練11 026次的total loss圖，損失率降低到了10%。

（6）通過(guò)TensorFlow Object Detection API庫(kù)的測(cè)試程序object＿detection＿tutorial.py對(duì)訓(xùn)練好的PB 模型文件進(jìn)行測(cè)試，考查模型訓(xùn)練的結(jié)果。圖7為使用藥片板數(shù)據(jù)集對(duì)模型訓(xùn)練11 026次后的識(shí)別效果圖，可看出已能夠識(shí)別出全部的藥片板。

圖5 模型訓(xùn)練9 817次的識(shí)別率

圖6 模型訓(xùn)練11 026次的整體損失率圖

圖7 模型訓(xùn)練11 026次后的識(shí)別效果圖

3.4 垃圾識(shí)別模型在TX2上的部署

在計(jì)算機(jī)上訓(xùn)練好的垃圾識(shí)別模型文件，需要在NVIDIA Jeston TX2（后面簡(jiǎn)稱TX2）平臺(tái)上進(jìn)行部署。TX2平臺(tái)上需事先安裝好NVIDIA的神經(jīng)網(wǎng)絡(luò)推理加速框架TensorRT。TensorRT可以將TensorFlow、Caffe、PyTorch和MXnet等框架的深度學(xué)習(xí)模型進(jìn)行解析，并針對(duì)NVIDIA設(shè)備進(jìn)行優(yōu)化和部署加速［14-15］。

實(shí)驗(yàn)流程：

（1）將訓(xùn)練好的最新PB模型文件、lable map文件和ssd＿mobilenet＿v1＿coco.config 模型配置文件拷貝到TensorRT的data目錄下。

（2）為增加對(duì)新模型文件的支持，需要修改TensorRT的utils目錄下的egohands＿models.py文件，添加對(duì)新模型的支持。

（3）打開(kāi)TensorRT 中的camera＿tf＿trt.py 文件，修改模型名稱DEFAULT＿M(jìn)ODEL和lable map文件路徑DEFAULT＿LABELMAP。

（4）修改TensorRT的utils目錄下的visualization.py文件，實(shí)現(xiàn)對(duì)識(shí)別結(jié)果的串口輸出。

經(jīng)過(guò)以上實(shí)驗(yàn)步驟，即可完成模型部署。圖8～11為對(duì)藥片板的識(shí)別結(jié)果，可見(jiàn)訓(xùn)練好的模型在不同光線、不同姿態(tài)、不同種類、遠(yuǎn)距離和背景雜亂情況下，都可以實(shí)現(xiàn)對(duì)藥片板的準(zhǔn)確識(shí)別。

圖8 強(qiáng)光環(huán)境下的識(shí)別結(jié)果

圖9 暗光、擺放姿態(tài)改變情況下的識(shí)別結(jié)果

圖10 暗光、不同種類藥片板識(shí)別結(jié)果

圖11 暗光、遠(yuǎn)距離、背景雜亂情況下識(shí)別結(jié)果

3.5 垃圾分揀機(jī)器人實(shí)驗(yàn)平臺(tái)

圖12 為垃圾分揀機(jī)器人實(shí)驗(yàn)平臺(tái)實(shí)物圖，其中包括：垃圾投放入口、垃圾分離傳輸帶、垃圾承載托盤(pán)、USB高清攝像頭、4軸機(jī)械臂、NVIDIA Jeston TX2、STM32控制板以及智能垃圾桶等設(shè)備。表1列出了目前實(shí)驗(yàn)平臺(tái)可以識(shí)別的、方便教學(xué)的垃圾種類。

圖12 垃圾分揀機(jī)器人實(shí)驗(yàn)平臺(tái)

表1 垃圾分揀機(jī)器人可識(shí)別的垃圾種類列表

4 結(jié) 語(yǔ)

本文以垃圾分揀為實(shí)際工作場(chǎng)景，基于TensorFlow深度學(xué)習(xí)框架，采用MobileNet SSD模型，通過(guò)NVIDIA Jetson TX2機(jī)器學(xué)習(xí)平臺(tái)和智能機(jī)械臂控制系統(tǒng)，實(shí)現(xiàn)了垃圾分揀機(jī)器人實(shí)驗(yàn)平臺(tái)的設(shè)計(jì)和研制，并進(jìn)行了相應(yīng)的垃圾分類實(shí)驗(yàn)設(shè)計(jì)；為人工智能相關(guān)專業(yè)的教學(xué)、科研和課外學(xué)術(shù)競(jìng)賽提供了一套完整的教學(xué)實(shí)驗(yàn)平臺(tái)，具有良好的實(shí)用價(jià)值。本垃圾分揀機(jī)器人在2019世界機(jī)器人大會(huì)工業(yè)機(jī)器人設(shè)計(jì)大賽中，獲得一等獎(jiǎng)，并獲得現(xiàn)場(chǎng)專家和觀眾的一致好評(píng)。