图像识别垃圾分类代码 图像识别垃圾分类程序

垃圾分类小程序

1. 结构

• 微信小程序
• 图像识别模型
• API

2. 微信小程序设计

对于小程序的界面设计，我设计的比较简单。如图：

当然。界面完全可以依照自己的想法进行设计。对ui设计我考虑的不多。主要还是在模型上重视一点。

代码：

<!--index.wxml-->
<view class="container">
  <image class='background' src="http://file05.16sucai.com/2015/0714/69b492e5d11a41e047aa47a090b11ebe.jpg" mode="aspectFill"></image>
  <view>
    <button class='btn1' bindtap="chooseimage">识别垃圾</button>
  </view>
  <modal hidden="{{hidden}}" confirm-text="确定" bindconfirm="confirm" no-cancel="{{nocancel}}">
      <view>
        <image class="image" src="{{tempFilePaths}}"></image>
      </view>
      <view class="result">
      {{result}}
      </view>
  </modal>
</view>

//index.js
//获取应用实例
const app = getApp()

Page({
  data: {
    tempFilePaths: '',
    hidden: true,
    nocancel: true,
    image_url:'',
    result:''
  },
  confirm: function () {
    this.setData({
      hidden: !this.data.hidden,
      result: ''
    });
  },
  chooseimage: function () {
    var that = this;
    wx.showActionSheet({
      itemList: ['从相册选择', '拍照'],
      itemColor: "#9978c4",
      success: function (res) {
        if (!res.cancel) {
          if (res.tapIndex == 0) {
            that.chooseWxImage('album')
          } else if (res.tapIndex == 1) {
            that.chooseWxImage('camera')
          }
        }
      }
    })
  },
  chooseWxImage: function (type) {
    var that = this;
    wx.chooseImage({
      sizeType: ['original', 'compressed'],
      sourceType: [type],
      success: function (res) {
        console.log(res);
        that.uploadImg(res.tempFilePaths[0])
        that.setData({
          tempFilePaths: res.tempFilePaths[0],
          hidden: false
        })
      },
    })
  },
  uploadImg: function (filePath){
    var that = this;
    wx.uploadFile({
      url: 'http://127.0.0.1:5000/classification',
      filePath: filePath,
      name: 'file',
      success: function(res){
        that.setData({
          result: unescape(res.data.replace(/\\/g, "%"))
        })
      }
    })
  }
})

3. API

接口，我采用的是flask-restful，我只写了一个用于识别的接口。

import flask
from flask_restful import Resource, Api
from flask import request, Flask
import tensorflow as tf
import numpy as np
from resnet import resnet18, resnet34
import tensorflow as tf
from tensorflow.keras import layers, Sequential, regularizers, optimizers
import tensorflow.keras as keras
import log
import traceback
from tensorflow.keras.applications.resnet50 import ResNet50
import json

app = Flask(__name__)
api = Api(app)


def load_image(image_path):
    image = tf.io.read_file(image_path)
    image = tf.image.decode_jpeg(image)
    image = tf.image.resize(image, [128, 128], method='nearest')
    return image

model = ResNet50(
    weights=None,
    classes=40,
    input_shape=(128, 128, 3)
)
model.trainable = False
model.load_weights('model/model.h5')

with open('model/garbage_classify_rule.json', 'r', encoding='utf-8') as f:
    dit = json.load(f)


def perdict(image_path):
    images = []
    images.append(load_image(image_path))
    images = np.array(2 * tf.cast(images, dtype=tf.float32) / 255. - 1, dtype = np.float32)
    
    res = model.predict(images)
    max = 0
    classifaction = 0
    for i in range(len(res[0])):
        if res[0][i] > max:
            classifaction = i
            max = res[0][i]
    return dit[str(classifaction)]



class Classification(Resource):
    def post(self):
        try:
            img = request.files.get('file')
            path = 'images/%s' % (img.filename)
            img.save(path)
            res = perdict(image_path=path)
            print(res)
            return str(res)
        except:
            log.logging.error('%s' % traceback.format_exc())
            return "1"

api.add_resource(Classification, '/classification')

if __name__ == "__main__":
    app.run(debug = True)

4.模型构建

这里才是整个项目的核心。

这里我采用的是tensorflow深度学习框架（v2.3）。网络模型就用的resnet。数据集使用的是华为垃圾分类数据集。

{
    "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": "有害垃圾/过期药物"
}

4.1 数据处理

首先需要对数据集中的标签以及图片进行读取。

def __build_img_label(self, txt_path, train_data_dir):
        try:
            res = []
            content = ''
            with open(txt_path, 'r', encoding='utf-8') as txt_f:
                content = txt_f.read()

            res = content.split(',')
            
            if len(res) == 2:
                res[0] = '%s/%s'% (train_data_dir, res[0].strip())
                res[1] = int(res[1].strip())
            
            return res
            
        except:
            log.logging.error('%s' % traceback.format_exc())

处理好路径以后就开始读取图片：

def __load_train_data(self):
        try:
            train_data = []
            img_txt_files = []
            for root, dirs, files in os.walk(self.__data_path):
                img_txt_files = files
            for _file in img_txt_files:
                if '.txt' in _file:
                    txt_file = os.path.join(self.__data_path, _file)
                    train_data.append(self.__build_img_label(txt_file, self.__data_path))

            train_images = []
            train_labels = []

            def load_image(image_path):
                image = tf.io.read_file(image_path)
                image = tf.image.decode_jpeg(image)
                image = tf.image.resize(image, [128, 128], method='nearest')
                return image

            num = 0
            for data in train_data:
                train_images.append(load_image(data[0]))
                train_labels.append([data[1]])
                num += 1
                if num % 1000 == 0:
                    print('图片加载：%s' % (num )) 
            x_train = np.zeros(shape = [num, 128, 128, 3], dtype = np.float32)
            for i in range(len(train_images)):
                x_train[i,:] = train_images[i]
            return x_train, np.array(train_labels)      
        except:
            log.logging.error('%s' % traceback.format_exc())

我这里设置的图片大小为128*128，对于图片维度，根据自己gpu进行选择。

全部代码（DataProcess.py）:

# coding: UTF-8
import json
import log
import traceback
import os
from tensorflow import keras
import tensorflow as tf
import gzip
import numpy as np

class DataProcess(object):
    def __init__(self, train_data_path, rule_path):
        self.__data_path = train_data_path
        self.__rule_path = rule_path

    def load_label(self):
        try:
            with open(self.__rule_path, 'r', encoding='utf-8') as load_f:
                label_rule = json.load(load_f)
            # print(label_rule)
            return load_f
        except:
            log.logging.error('%s' % traceback.format_exc())

    def build_tensor_data(self):
        # 构建训练集
        train_images, train_labels = self.__load_train_data()

        return train_images, train_labels
    

    def __load_train_data(self):
        try:
            train_data = []
            img_txt_files = []
            for root, dirs, files in os.walk(self.__data_path):
                img_txt_files = files
            for _file in img_txt_files:
                if '.txt' in _file:
                    txt_file = os.path.join(self.__data_path, _file)
                    train_data.append(self.__build_img_label(txt_file, self.__data_path))

            train_images = []
            train_labels = []

            def load_image(image_path):
                image = tf.io.read_file(image_path)
                image = tf.image.decode_jpeg(image)
                image = tf.image.resize(image, [128, 128], method='nearest')
                return image

            num = 0
            for data in train_data:
                train_images.append(load_image(data[0]))
                train_labels.append([data[1]])
                num += 1
                if num % 1000 == 0:
                    print('图片加载：%s' % (num )) 
            x_train = np.zeros(shape = [num, 128, 128, 3], dtype = np.float32)
            for i in range(len(train_images)):
                x_train[i,:] = train_images[i]
            return x_train, np.array(train_labels)      
        except:
            log.logging.error('%s' % traceback.format_exc())
    

    def __build_img_label(self, txt_path, train_data_dir):
        try:
            res = []
            content = ''
            with open(txt_path, 'r', encoding='utf-8') as txt_f:
                content = txt_f.read()

            res = content.split(',')
            
            if len(res) == 2:
                res[0] = '%s/%s'% (train_data_dir, res[0].strip())
                res[1] = int(res[1].strip())
            
            return res
            
        except:
            log.logging.error('%s' % traceback.format_exc())


if __name__ == "__main__":
    data_process = DataProcess('garbage_classify/train_data', 'garbage_classify/garbage_classify_rule.json')
    train_images, train_labels = data_process.build_tensor_data()

    np.save('train_images.npy', train_images)
    np.save('train_labels.npy', train_labels)

    print(train_images.shape)

4.2 网络模型

然后是resnet网络模型（resnet.py）：

import tensorflow as tf
from tensorflow.keras import layers, Sequential
from tensorflow import keras


class BasicBlock(layers.Layer):

    def __init__(self, filter_num, stride = 1):
        super(BasicBlock, self).__init__()
        self.conv1 = layers.Conv2D(filter_num, (3, 3), strides=stride, padding='same')
        self.bn1 = layers.BatchNormalization()
        self.relu = layers.Activation('relu')

        self.conv2 = layers.Conv2D(filter_num, (3, 3), strides = 1, padding='same')

        self.bn2 = layers.BatchNormalization()

        if stride != 1:
            self.downsample = Sequential()
            self.downsample.add(layers.Conv2D(filter_num, (1, 1), strides = stride))

        else:
            self.downsample = lambda x: x

    
    def call(self, inputs, training = None):
        out = self.conv1(inputs)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)

        identity = self.downsample(inputs)

        output = layers.add([out, identity])

        output = tf.nn.relu(output)

        return output


class ResNet(keras.Model):
    def __init__(self, layer_dims, num_classes = 10):
        super(ResNet, self).__init__()
        self.stem = Sequential([
            layers.Conv2D(64, (3, 3), strides = (1, 1)),
            layers.BatchNormalization(),
            layers.Activation('relu'),
            layers.MaxPool2D(pool_size = (2, 2), strides = (1, 1), padding = 'same')
        ])

        self.layer1 = self.build_resblock(64, layer_dims[0])
        self.layer2 = self.build_resblock(128, layer_dims[1])
        self.layer3 = self.build_resblock(256, layer_dims[2])
        self.layer4 = self.build_resblock(512, layer_dims[3])

        self.avgpool = layers.GlobalAveragePooling2D()

        self.fc = layers.Dense(num_classes)

        
    def call(self, inputs, training = None):
        x = self.stem(inputs)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        
        x = self.avgpool(x)
        x = self.fc(x)

        return x

    def build_resblock(self, filter_num, blocks, stride = 2):
        res_block = Sequential()
        res_block.add(BasicBlock(filter_num, stride))

        for _ in range(1, blocks):
            res_block.add(BasicBlock(filter_num, stride = 1))
        
        return res_block


def resnet18():
    return ResNet([2, 2, 2, 2], num_classes = 40)


def resnet34():
    return ResNet([3, 4, 6, 3], num_classes = 40)

这里代码完全参照的《TensorFlow深度学习》。

关于这个模型的相关解释，大家可以另行搜索。资料很多，而且都容易理解。很多博主解释得很详细。

4.3 训练

train.py

import tensorflow as tf
from tensorflow.keras import layers, Sequential, regularizers, optimizers
import tensorflow.keras as keras
from DataProcess import DataProcess
from resnet import resnet18, resnet34
import numpy as np
import log
from rnnnet import create_model

data_process = DataProcess('garbage_classify/train_data', 'garbage_classify/garbage_classify_rule.json')
train_images, train_labels = data_process.build_tensor_data()

train_labels = tf.squeeze(train_labels, axis=1)

def preprocess(x, y):
    # 将数据映射到-1~1
    x = 2*tf.cast(x, dtype=tf.float32) / 255. - 1
    y = tf.cast(y, dtype=tf.int32) # 类型转换
    return x,y

train_db = tf.data.Dataset.from_tensor_slices((train_images,train_labels))
train_db = train_db.shuffle(15000).map(preprocess).batch(256)

model = resnet18()
model.trainable = False

model.build(input_shape = train_images.shape)

model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer=keras.optimizers.SGD(learning_rate=0.1, momentum=0.9),
    metrics=['accuracy'])

model.fit(train_db,
        epochs=25,
        )


model.save_weights('model/model.h5')

这里是采用18层的网络结构。

我实际项目中采用的是50层的，直接引用的python的内置包。

代码如下：

import tensorflow as tf
from tensorflow.keras import layers, Sequential, regularizers, optimizers
import tensorflow.keras as keras
import numpy as np
from tensorflow.keras.applications.resnet50 import ResNet50


train_images = np.load('train_images.npy')
train_labels = np.load('train_labels.npy')
# print(train_images.shape)

train_labels = tf.squeeze(train_labels, axis=1)

def preprocess(x, y):
    # 将数据映射到-1~1
    x = 2*tf.cast(x, dtype=tf.float32) / 255. - 1
    y = tf.cast(y, dtype=tf.int32) # 类型转换
    return x,y

train_db = tf.data.Dataset.from_tensor_slices((train_images,train_labels))
train_db = train_db.shuffle(15000).map(preprocess).batch(256)

model = ResNet50(
    weights=None,
    classes=40,
    input_shape=(128, 128, 3)
)


model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer='adam',
    metrics=['accuracy'])


model.summary() 


model.fit(train_db,
    epochs=25)

model.save_weights('model.h5')

训练过程：

总结一下：
怎么项目的重点主要在于模型的选择以及数据的处理，这也是导致整个模型识别率的因素。对于数据处理更是重中之重。