【深度学习实践】垃圾检测

简介

本项目使用深度学习目标检测开源框架PaddleDetection中的yolox算法实现了垃圾检测，本文包含了从头训练yolox模型和直接使用训练好的模型进行推理的代码及相关权重。

一、数据集准备

本次训练的数据集为coco格式，共包含150张垃圾的照片，如下图所示：

数据集链接：传送门

二、准备PaddleDetection

在训练yolox之前，先准备PaddleDetection环境。

1、下载PaddleDetection

从github下载PaddleDetection：

git clone https://github.com/PaddlePaddle/PaddleDetection.git

2、安装依赖

安装依赖包

cd PaddleDetection
pip install -r requirements.txt

3、安装PaddleDetection

python setup.py install

三、训练YOLOX

提示：使用作者提供的模型进行垃圾检测推理请跳过此步骤。

准备好数据集和PaddleDetection之后，就可以训练模型了。

1、准备配置文件

为了防止数据集路径配置错误，建议将垃圾数据集路径配置为：PaddleDetection/dataset/trash_dataset，该目录下包含Annotations和Images两个文件夹，Images里面有150个垃圾图片，Annotations里面有json标注文件。PaddleDetection/dataset目录结构如下：

数据集位置确定以后，打开PaddleDetection/configs/dataset/coco_detection.yml，修改dataset_dir设置，修改为自己的数据集路径，如果使用我一样的数据集路径，那么配置文件的内容如下：

metric: COCO
num_classes: 1

TrainDataset:
  !COCODataSet
    image_dir: Images
    anno_path: Annotations/train.json
    dataset_dir: dataset/trash_dataset
    data_fields: ['image', 'gt_bbox', 'gt_class', 'is_crowd']

EvalDataset:
  !COCODataSet
    image_dir: Images
    anno_path: Annotations/val.json
    dataset_dir: dataset/trash_dataset

TestDataset:
  !ImageFolder
    image_dir: Images
    anno_path: Annotations/test.json # also support txt (like VOC's label_list.txt)
    dataset_dir: dataset/trash_dataset # if set, anno_path will be 'dataset_dir/anno_path'

2、训练YOLOX

PaddleDetection/configs/yolox目录下有多个yolox模型的配置文件，可以根据自己的需要选择合适的模型，这里使用yolox_tiny_300e_coco.yml配置文件进行训练，相关参数可以自行修改，命令如下（在PaddleDetection目录下）：

python tools/train.py -c configs/yolox/yolox_tiny_300e_coco.yml 

3、验证

训练好了模型以后，需要验证训练的模型精度，命令如下（weights指向训练好的权重路径）：

python tools/eval.py -c configs/yolox/yolox_tiny_300e_coco.yml -o weights=output/yolox_tiny_300e_coco/model_final.pdparams 

四、模型预测

本章节使用训练好的模型进行推理，命令如下（在PaddleDetection目录下运行）：

python tools/infer.py -c configs/yolox/yolox_tiny_300e_coco.yml --infer_img=testImage/0f16-hzmafvm6151836.jpg -o weight=output/yolox_tiny_300e_coco/model_final.pdparams

本文代码、数据集、测试图片、训练好的权重都可以在链接【使用PaddleDetection实现垃圾检测】下载，按照【二、准备PaddleDetection】配置好环境即可直接运行。

原图为：

预测结果为：

五、参考

https://aistudio.baidu.com/aistudio/projectdetail/3846170?channelType=0&channel=0