带你动手做AI版的垃圾分类

摘要：本案例将使用YOLOX模型，实现一个简单的垃圾分类应用。

本文分享自华为云社区《ModelBox社区案例 - 使用YOLOX做垃圾分类》，作者：HWCloudAI。

1 ModelBox社区案例 - 使用YOLOX做垃圾分类

本案例将使用YOLOX模型，实现一个简单的垃圾分类应用，最终效果如下所示：

本案例所需资源（代码、模型、测试数据等）均可从garbage_det下载

1.1 模型训练与转换

模型采用的是YOLOX网络结构，YOLOX是YOLO系列的优化版本，引入了解耦头、数据增强、无锚点以及标签分类等目标检测领域的优秀进展，拥有较好的精度表现，同时对工程部署友好。训练使用的是“华为云杯”生活垃圾图片分类数据集，该数据集包含一次性快餐盒、果皮果肉、旧衣服等44个类别，共14964张图片。其中，训练集与验证集划分比例为4/1，下图为模型迭代个300个Epoch取得的结果：

ModelArts提供了包括数据标注，训练环境，预置算法在内的丰富的功能，甚至可以通过订阅预置算法实现0代码的模型训练工作。当然你也可以在本地训练自己的模型。我们假设你现在已经拥有了训练好的模型，接下来我们需要将训练好的模型转换成为可以在开发板上运行的模型。我们发布了开发板模型转换案例，参见RK3568模型转换验证案例：

在这个案例中我们演示了从环境适配到模型的转换验证的全流程样例代码，开发者可以通过“Run in ModelArts”一键将Notebook案例在ModelArts控制台快速打开、运行以及进行二次开发等操作。

1.2 应用开发

打开VS Code，连接到ModelBox sdk所在目录或者远程开发板，开始进行垃圾分类应用的开发。下面以RK3568版本为例进行说明，其他版本与之类似。

1.2.1 1）下载模板

执行python solution.py -l可看到当前公开的技能模板：

███ $ python solution.py -l

...

Solutions name:

mask_det_yolo3

…

hand_det_yolox

hand_tracking_yolox

single_hand_pose_yolox_mbv2

multi_hand_pose_yolox_mbv2

结果中的hand_det_yolox即为手部检测应用模板，可使用如下命令下载模板：

███ $ python solution.py -s hand_det_yolox

...

solution.py工具的参数中，-l 代表list，即列出当前已有的模板名称；-s 代表solution-name，即下载对应名称的模板。下载下来的模板资源，将存放在ModelBox核心库的solution目录下。

1.2.2 2）创建工程

在ModelBox sdk目录下使用create.py创建garbage_det工程，末尾-s参数，表示将使用后面参数值代表的模板创建工程，而不是创建空的工程。

███/modelbox$ python create.py -t server -n garbage_det -s hand_det_yolox

sdk version is modelbox-xxx

success: create garbage_det in ███/modelbox/workspace

workspace目录下将创建出garbage_det工程，工程内容如下所示：

garbage_det

|--bin

│  |--main.bat：应用执行入口

│  |--mock_task.toml：应用在本地执行时的输入输出配置，此应用默认使用本地视频文件为输入源，最终结果输出到另一本地视频文件，可根据需要修改

|--CMake：存放一些自定义CMake函数

|--data：存放应用运行所需要的图片、视频、文本、配置等数据

│  |--hand.mp4：手部检测测试用视频文件—>替换为自己的视频

|--dependence

│  |--modelbox_requirements.txt：应用运行依赖的外部库在此文件定义

|--etc

│  |--flowunit：应用所需的功能单元存放在此目录

│  │  |--cpp：存放C++功能单元编译后的动态链接库，此应用没有C++功能单元

│  │  |--yolox_post：手部检测使用的是YOLOX模型，此处即为后处理功能单元（修改toml文件的类别参数和py文件的draw函数）

|--flowunit_cpp：存放C++功能单元的源代码，此应用没有C++功能单元

|--graph：存放流程图

│  |--garbage_det.toml：默认流程图，使用本地视频文件作为输入源

│  |--garbage_det_camera.toml：摄像头输入对应的流程图

│  |--modelbox.conf：modelbox相关配置

|--hilens_data_dir：存放应用输出的结果文件、日志、性能统计信息

|--model：推理功能单元目录

│  |--detect_hand：手部检测推理功能单元

│  │  |--detect_hand.toml：手部检测推理功能单元的配置文件

│  │  |--yolox_hand.onnx：手部检测onnx模型—>更改为自己的模型

|--build_project.sh：应用构建脚本

|--CMakeLists.txt

|--rpm：打包rpm时生成的目录，将存放rpm包所需数据

|--rpm_copyothers.sh：rpm打包时的辅助脚本

1.2.3 3）修改后处理功能单元 yolox_post

a. 修改yolox_post.toml流程图，将其内容修改为（以Windows版ModelBox为例）：

b. 修改yolox_post.py的draw函数实现如下：

def draw(self, img, bboxes):

        h, w, c = img.shape

        thickness = 2

        font_scale = 1

        text_font = cv2.FONT_HERSHEY_SIMPLEX

        clss_to_text = {

        0: "Disposable snack box",

        1: "Books and papers",

        2: "Power bank",

        3: "Leftovers",

        4: "Package",

        5: "Trash can",

        6: "Plastic utensils",

        7: "Plastic toys",

        8: "Plastic coat hanger",

        9: "Big Bones",

        10: "Dry battery",

        11: "Express paper bag",

        12: "Plug wire",

        13: "Old clothes",

        14: "The can",

        15: "Pillow",

        16: "Skin and pulp",

        17: "Stuffed animal",

        18: "Defacing plastic",

        19: "Soiled paper",

        20: "Toiletries",

        21: "Cigarette butts",

        22: "Toothpick",

        23: "Glassware",

        24: "Block",

        25: "Chopsticks",

        26: "Carton carton",

        27: "Pot",

        27: "Tea residue",

        29: "Vegetable help vegetable leaf",

        30: "Shell",

        31: "The spice bottle",

        32: "Paste",

        33: "Expired drugs",

        34: "Bottle",

        35: "Metal kitchenware",

        36: "Metal ware",

        37: "Metal food cans",

        38: "Pot",

        39: "Ceramic vessels",

        40: "Shoes",

        41: "Edible oil drum",

        42: "Beverage bottle",

        43: "Bones"

        }

        for box in bboxes:

            x1, y1, x2, y2, score, clss = box

            cv2.putText(img, clss_to_text[int(clss)]+': '+"{:.3}".format(score*100)+'%', (int(x1 * w)+10, int(y1 * h)+30),text_font, font_scale, (0, int(clss+1)*5, 0), thickness)

            cv2.rectangle(img, (int(x1 * w), int(y1 * h)), (int(x2 * w), int(y2 * h)), (0, int(clss+1)*5, 0), 3)

1.2.4 4）修改输入输出配置

我们需要准备一个mp4文件拷贝到data文件夹下，我们使用测试视频garbage.mp4，然后打开工程目录下bin/mock_task.toml文件，修改其中的任务输入和任务输出配置为如下内容：

[input]

type = "url"

url = "../data/garbage.mp4"

[output]

type = “local”

url = “…/hilens_data_dir/garbage_detection_result.mp4”

该流程图在本地运行时的逻辑过程是：data_source_parser解析bin/mock_task.toml文件中输入配置的data/garbage.mp4文件，video_demuxer和video_decoder对该文件进行解码，resize、packed_planar_transpose、normalize对原始图像进行缩放、转码、归一化等预处理，然后detect_garbage在预处理后的图像上进行垃圾检测，yolox_post从推理结果中解码出检测框，并把检测框画到原始图像上，最后video_out将图像输出到bin/mock_task.toml文件中输出配置的hilens_data_dir/garbage_detection_result.mp4文件中。

1.2.5 5）用启动脚本执行应用

启动应用前执行build_project.sh进行工程构建，该脚本将编译自定义的C++功能单元（本应用不涉及）、将应用运行时会用到的配置文件转码为Unix格式（防止执行过程中的格式错误）：

███$ ./build_project.sh

dos2unix: converting file xxx.toml to Unix format...

...

build success: you can run main.bat in ./bin folder

Press ‘p’ to pause…, any key to exit

然后执行bin/main.bat运行应用：

███$ ./bin/main.bat

…

运行结束后在hilens_data_dir目录下生成了garbage_detection_result.mp4文件，可以打开查看：

1.2.6 6）用摄像头检测

打开工程目录下bin/mock_task.toml文件，修改其中的任务输入和任务输出配置为如下内容：

[input]

type = "url"

url = "0"  # 表示0号摄像头，即PC自带摄像头，若PC无摄像头需外接USB摄像头

[output]

type = “local”

url = “0:garbage_det” # 表示名为garbage_det的本地窗口

即使用编号为0的摄像头（默认为PC自带的摄像头），输出画面显示到名为garbage_det的本地屏幕窗口中。

1.2.7 7）运行应用

执行bin/main.bat camera运行应用，将会自动弹出实时的垃圾分类检测画面：

1.3 打包部署

1.3.1 打包

调试完成后，同样可以通过create.py脚本将应用打包发布：

python ./create.py -t rpm -n garbage_det

控制台中输出：

sdk version is modelbox-win10-x64-1.1.0.5

call mb-pkg-tool pack [folder] > [rpm file] to building rpm, waiting...

success: create garbage_det.rpm in D:\modelbox-win10-x64-1.1.0.5/workspace/garbage_det

等待稍许，可以看到项目工程下已经生成了rpm文件夹和打包好的应用：