基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统

今天在查资料的时候在网上看到一篇文章,博主是基于TensorFlow实现的CNN来完成对垃圾分类识别的,想到最近正好在使用Pytorch就想也做一下,就当是项目开发实践了。先看下动态操作效果:

原文在这里,截图如下所示:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第1张图片

 感兴趣的话可以自行移步原文。

原文实现的垃圾分类系统更为详尽,一共细分了40个不同的子类,如下所示:

{
    "0": "其他垃圾/一次性快餐盒",
    "1": "其他垃圾/污损塑料",
    "2": "其他垃圾/烟蒂",
    "3": "其他垃圾/牙签",
    "4": "其他垃圾/破碎花盆及碟碗",
    "5": "其他垃圾/竹筷",
    "6": "厨余垃圾/剩饭剩菜",
    "7": "厨余垃圾/大骨头",
    "8": "厨余垃圾/水果果皮",
    "9": "厨余垃圾/水果果肉",
    "10": "厨余垃圾/茶叶渣",
    "11": "厨余垃圾/菜叶菜根",
    "12": "厨余垃圾/蛋壳",
    "13": "厨余垃圾/鱼骨",
    "14": "可回收物/充电宝",
    "15": "可回收物/包",
    "16": "可回收物/化妆品瓶",
    "17": "可回收物/塑料玩具",
    "18": "可回收物/塑料碗盆",
    "19": "可回收物/塑料衣架",
    "20": "可回收物/快递纸袋",
    "21": "可回收物/插头电线",
    "22": "可回收物/旧衣服",
    "23": "可回收物/易拉罐",
    "24": "可回收物/枕头",
    "25": "可回收物/毛绒玩具",
    "26": "可回收物/洗发水瓶",
    "27": "可回收物/玻璃杯",
    "28": "可回收物/皮鞋",
    "29": "可回收物/砧板",
    "30": "可回收物/纸板箱",
    "31": "可回收物/调料瓶",
    "32": "可回收物/酒瓶",
    "33": "可回收物/金属食品罐",
    "34": "可回收物/锅",
    "35": "可回收物/食用油桶",
    "36": "可回收物/饮料瓶",
    "37": "有害垃圾/干电池",
    "38": "有害垃圾/软膏",
    "39": "有害垃圾/过期药物"
}

我这里没有那么多的细分的数据,从数据库中找到了一个四类别的垃圾识别数据集,主要包括:

有害垃圾
可回收垃圾
厨余垃圾
其他垃圾

本质都是一样的,都是属于多分类任务。所以这里不需要过于关注数据集类别的问题。

首先看下效果:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统

 经过一下午的开发,终于完结,项目整体截图如下所示:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第2张图片

 其中:

train.py负责整套模型的训练

项目底层使用的是mobilenet系列轻量级的识别模型,对于设备端的压力会更小一点,模型整体体积在8MB左右,一般的PC机在CPU小都是可以轻松推理的。

四种类型垃圾数据截图如下所示:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第3张图片

 厨余垃圾:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第4张图片

 有害垃圾:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第5张图片

 其他垃圾:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第6张图片

 可回收垃圾:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第7张图片

 results/目录用于存放模型的结果输出,如下所示:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第8张图片

 其中evaluation/目录是由evaluation.py模块对测试集进行评估计算得到的评估结果:

我保存了GPU服务器训练时的截图,如下:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第9张图片

 红框中的就是我们的模型,可以看到显存占用量不到4GB,还是不错的。

训练日志输出如下所示:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第10张图片

 Precision指标计算如下:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第11张图片

 Recall指标计算如下:

基于Pytorch框架的轻量级卷积神经网络垃圾分类识别系统_第12张图片

 记录一下,持续学习中,如有需要的话界面顶部有详情,可详细咨询。

你可能感兴趣的:(pytorch,cnn,分类)