简介

本项目使用深度学习目标检测开源框架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

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐