利用AI技术实现城市垃圾的自动分类

上海市作为我国的一线大城市,率先打响了垃圾分类的第一枪,截至到目前为止,垃圾分类工作已开展数月,由于不少人之前对于垃圾分类都缺乏了解,所以导致许多人对于垃圾分类工作很是苦恼。

图1 各类垃圾桶
图2 协助垃圾分类的志愿者

为了方便人们进行垃圾分类,现如今最常用的做法是在各个垃圾收放点都安排了许多的志愿者协助大家进行垃圾分类,这无疑大大地增加了人工劳动力。针对此问题,本文将AI技术用于了城市垃圾分类,只需对垃圾进行拍照就可自动识别出它的类别,对于还不了解垃圾分类的童鞋来说,简直不要太方便。


基于深度学习的城市垃圾分类

数据集介绍

本文使用的数据集为Gary Thung制作的公开图像数据集,数据集下载地址为:
https://github.com/garythung/trashnet/blob/master/data/dataset-resized.zip

实验数据集一共包括五大类:

  • 硬纸板cardboard
  • 玻璃glass
  • 金属metal
  • 纸paper
  • 塑料plastic

数据集图像尺寸均为512*384像素大小。按照3:1:1的比例将数据集划分为训练集、验证集和测试集,即每个类别都包含243张训练集、80张验证集和80张测试集。由于数据集数量有限,为了获得更好的识别效果,文中还利用180度旋转和镜像的方式对训练集进行了扩充,通过数据扩充训练集数据量变为了原先的4倍,也就是说每个类别都包含972张训练集图像。

图3 数据集图像示例

模型训练

文中构建了一个八层的卷积神经网络模型,网络结构图如下所示:

图4 网络结构图

文中模型的详细参数如下所示:

图5 网络模型具体参数

本文在Tensorflow框架下完成了网络搭建,网络模型代码如下:

def Mynet(image, keepprob=0.5):
    # 定义卷积层1,卷积核大小,偏置量等各项参数参考下面的程序代码,下同。
    with tf.name_scope("conv1") as scope:
        kernel = tf.Variable(tf.truncated_normal([11, 11, 3, 64], dtype=tf.float32, stddev=1e-1, name="weights"))
        conv = tf.nn.conv2d(image, kernel, [1, 4, 4, 1], padding="SAME")
        biases = tf.Variable(tf.constant(0.0, dtype=tf.float32, shape=[64]), trainable=True, name="biases")
        bias = tf.nn.bias_add(conv, biases)
        conv1 = tf.nn.relu(bias, name=scope)

        pass

    # LRN层
    lrn1 = tf.nn.lrn(conv1, 4, bias=1.0, alpha=0.001/9, beta=0.75, name="lrn1")

    # 最大池化层
    pool1 = tf.nn.max_pool(lrn1, ksize=[1,3,3,1], strides=[1,2,2,1],padding="VALID", name="pool1")

    # 定义卷积层2
    with tf.name_scope("conv2") as scope:
        kernel = tf.Variable(tf.truncated_normal([5,5,64,192], dtype=tf.float32, stddev=1e-1, name="weights"))
        conv = tf.nn.conv2d(pool1, kernel, [1, 1, 1, 1], padding="SAME")
        biases = tf.Variable(tf.constant(0.0, dtype=tf.float32, shape=[192]), trainable=True, name="biases")
        bias = tf.nn.bias_add(conv, biases)
        conv2 = tf.nn.relu(bias, name=scope)
        pass

    # LRN层
    lrn2 = tf.nn.lrn(conv2, 4, bias=1.0, alpha=0.001 / 9, beta=0.75, name="lrn2")

    # 最大池化层
    pool2 = tf.nn.max_pool(lrn2, ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1], padding="VALID", name="pool2")

    # 定义卷积层3
    with tf.name_scope("conv3") as scope:
        kernel = tf.Variable(tf.truncated_normal([3,3,192,384], dtype=tf.float32, stddev=1e-1, name="weights"))
        conv = tf.nn.conv2d(pool2, kernel, [1, 1, 1, 1], padding="SAME")
        biases = tf.Variable(tf.constant(0.0, dtype=tf.float32, shape=[384]), trainable=True, name="biases")
        bias = tf.nn.bias_add(conv, biases)
        conv3 = tf.nn.relu(bias, name=scope)
        pass

    # 定义卷积层4
    with tf.name_scope("conv4") as scope:
        kernel = tf.Variable(tf.truncated_normal([3,3,384,256], dtype=tf.float32, stddev=1e-1, name="weights"))
        conv = tf.nn.conv2d(conv3, kernel, [1, 1, 1, 1], padding="SAME")
        biases = tf.Variable(tf.constant(0.0, dtype=tf.float32, shape=[256]), trainable=True, name="biases")
        bias = tf.nn.bias_add(conv, biases)
        conv4 = tf.nn.relu(bias, name=scope)
        pass

    # 定义卷积层5
    with tf.name_scope("conv5") as scope:
        kernel = tf.Variable(tf.truncated_normal([3,3,256,256], dtype=tf.float32, stddev=1e-1, name="weights"))
        conv = tf.nn.conv2d(conv4, kernel, [1, 1, 1, 1], padding="SAME")
        biases = tf.Variable(tf.constant(0.0, dtype=tf.float32, shape=[256]), trainable=True, name="biases")
        bias = tf.nn.bias_add(conv, biases)
        conv5 = tf.nn.relu(bias, name=scope)
        pass

    # 最大池化层
    pool5 = tf.nn.max_pool(conv5, ksize=[1,3,3,1], strides=[1,2,2,1], padding="VALID", name="pool5")

    # 全连接层
    flatten = tf.reshape(pool5, [-1, 15*11*256])

    weight1 = tf.Variable(tf.truncated_normal([15*11*256, 4096], mean=0, stddev=0.01))

    fc1 = tf.nn.sigmoid(tf.matmul(flatten, weight1))

    dropout1 = tf.nn.dropout(fc1, keepprob)

    weight2 = tf.Variable(tf.truncated_normal([4096, 4096], mean=0, stddev=0.01))

    fc2 = tf.nn.sigmoid(tf.matmul(dropout1, weight2))

    dropout2 = tf.nn.dropout(fc2, keepprob)

    weight3 = tf.Variable(tf.truncated_normal([4096, 5], mean=0, stddev=0.01))

    fc3 = tf.nn.sigmoid(tf.matmul(dropout2, weight3))

    return fc3

利用训练集对模型进行训练,共训练了200个epoch,如图5所示为网络训练结果。

图6 网络训练结果

模型测试

利用测试集对网络模型进行精度测试,最终的测试精度为77.75%,本文还增添了模型预测结果可视化显示功能,如下所示为模型的测试结果示例图,模型将预测结果在图片上方用英文名称显示。

可以看出,本文所搭建的模型对于大部分城市垃圾都可以做出正确分类识别,表明了本文的方法的可行性。

后记

本文所搭建的模型测试精度目前只有77.75%,精度仍然还有很大的上升空间,通过分析我认为主要以下几点原因:

  • 在实际生活中,城市垃圾的种类是十分繁多的,而本文实验中的数据量较少,尽管后续做了数据增强处理,使得训练集数据量达到了4860张,看起来这一数字比较可观,但是对于深度学习模型来说,数据量仍然过少,尤其是对于数据比较复杂的城市垃圾来说,这一数据量是难远远不够的。
  • 不同的数据集在不同的模型上表现效果也是不一致的,在后续,本文所搭建的深度学习网络模型还亟待进一步的优化。

本人知识有限,欢迎各位感兴趣的同仁提出宝贵意见和共同探讨,本人微信号:liuwei646102318,添加的时候请注明是从加的,希望共同成长。

你可能感兴趣的:(利用AI技术实现城市垃圾的自动分类)