毕业设计——基于 Pytorch的垃圾识别与分类系统的设计与实现

完整源码获取可私信博主

毕业设计——基于 Pytorch的垃圾识别与分类系统的设计与实现

随着全球范围内垃圾处理问题的日益严重，垃圾识别与分类技术成为了解决这一问题的关键手段。基于深度学习的垃圾识别与分类系统，尤其是利用PyTorch这样的强大框架，已经展现出了显著的成效。本文将对该系统的设计与实现进行综述，介绍其基本原理、关键技术、以及应用前景。

一、基本原理

基于PyTorch的垃圾识别与分类系统主要利用卷积神经网络（CNN）进行图像特征的提取和分类。CNN通过多层卷积、池化和全连接操作，能够从输入的垃圾图像中自动学习并提取出有用的特征信息。这些特征信息随后被用于训练分类器，以实现对垃圾图像的准确分类。

二、关键技术

数据预处理：由于垃圾图像数据通常具有多样性、复杂性和噪声等特点，因此需要对数据进行预处理以提高模型的性能。这包括图像缩放、归一化、增强等操作，以及标签的编码和转换。
模型构建：选择合适的CNN结构是系统设计的关键。常用的CNN结构包括AlexNet、VGGNet、ResNet等。这些结构具有不同的深度和复杂度，可以根据具体任务的需求进行选择和调整。
训练与优化：使用PyTorch框架进行模型的训练和优化。这包括选择合适的损失函数、优化器、学习率等超参数，以及进行模型的迭代训练和调整。同时，还需要利用验证集进行模型性能的评估，以防止过拟合现象的发生。
部署与测试：将训练好的模型部署到实际应用中，并对新的垃圾图像进行分类测试。这需要对模型进行前向传播计算，并输出相应的分类结果。

三、应用前景

基于PyTorch的垃圾识别与分类系统具有广泛的应用前景。首先，它可以应用于城市垃圾分类处理领域，帮助提高垃圾处理的效率和准确性。其次，它还可以应用于环保教育和宣传领域，通过展示不同垃圾的分类结果，提高公众对垃圾分类的认识和重视程度。此外，随着技术的不断发展，该系统还可以进一步拓展到智能家居、智能垃圾桶等领域，实现更加智能化和自动化的垃圾处理。

四、总结与展望

基于PyTorch的垃圾识别与分类系统是一种高效、准确的垃圾处理方法。通过利用深度学习和计算机视觉技术，该系统能够自动学习和提取垃圾图像的特征信息，并实现对垃圾的快速分类。然而，目前该系统仍面临一些挑战，如数据集的局限性、模型的泛化能力等问题。未来研究可以进一步探索如何优化模型结构、提高分类性能，并拓展其在更多领域的应用。

综上所述，基于PyTorch的垃圾识别与分类系统是一个具有重要意义的研究方向。随着技术的不断进步和应用场景的拓展，相信该系统将在未来发挥更大的作用，为解决全球垃圾处理问题提供有力支持。

源程序结构
在这里插入图片描述main.py: 模型训练代码
inference.py：前向传播
config.py: 程序运行时的相关参数
data/TrashSet.py：垃圾数据集类
images/文件夹：提供一些测试图片
models/mbv3_small.pth.tar：mobilenet v3 在 imagenet 预训练权重
models/mobilenetv3.py: 模型文件
utils/文件夹：提供一些关于 imagenet 的可用函数
模型训练结果
数据集
自己采集的垃圾分类数据集 TrashBig, 其中有 12 个小类别：
[‘bananapeel’, ‘battery’, ‘cardboard’, ‘clothes’, ‘drugs’, ‘glass’, ‘lightbulb’, ‘metal’, ‘paper’, ‘papercup’, ‘plastic’, ‘vegetable’]

这12个小类别分别属于四个大类别：
{‘Recyclables’: 2624, ‘Kitchen waste’: 939, ‘Hazardous waste’: 1581, ‘Other’: 963}

训练集图片共计 6000 张左右。验证数据集每类 60 张左右，共计 720 张左右。
训练结果
在这里插入图片描述
在这里插入图片描述大类别的分类精度都在 90% 以上。

使用自己的数据集训练模型
程序运行环境：Pytorch 1.0 以上版本、opencv-python、numpy、tensorboard。
数据集的组织要求：数据集目录下分为两个小目录：train 和 val。每个小目录下以类别名称命名的文件夹，存储对应类别的图片。
config.py 文件的修改：对 config.py 文件进行修改：

运行 main.py。
运行前向传播测试
直接运行 inference.py 即可测试。修改 config.py 中 DefaultConfig.InferWithGPU 参数即可切换前向传播使用 GPU 还是 CPU。对 inference.py 文件稍加修改即可对自己的图片进行分类。

主函数main源码：

import torchvision.models as models
from torchvision import transforms
from torch import nn
import torch
import cv2
from torch.utils.tensorboard import SummaryWriter
import time
import numpy as np

from config import DefaultConfig
from models.mobilenetv3 import MobileNetV3_Small
from data.TrashSet import TrashSet


def get_default_device():
    if torch.cuda.is_available():
     return torch.device('cuda')
    else:
        return torch.device('cpu')

def MyDevice(gpu=True):
    if gpu:
        return torch.device('cuda')
    else:
        return torch.device('cpu')

def train_with_valid(trashSet):
    dataloaders = trashSet.dataloaders
    model = MobileNetV3_Small(trashSet.class_num)
    device = torch.device('cpu')
    #device = MyDevice(DefaultConfig.TrainWithGPU)
    ckpt = torch.load("models/mbv3_small.pth.tar",map_location=torch.device('cpu'))
    model.load_state_dict({k.replace('module.', ''): v for k, v in ckpt['state_dict'].items()}, strict=False)
    model.to(device)

    for param in model.freeze_params():
        param.requires_grad = False
    for param in model.fine_tune_params():
        param.requires_grad = True

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.fine_tune_params(), lr=DefaultConfig.LearningRate, weight_decay=2.5e-5)
    exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=DefaultConfig.LRScheduleStep, gamma=0.1)

    startEpoch = 0
    writer = SummaryWriter(DefaultConfig.LogDir)
    class_acc_dict = dict()
    group_acc_dict = dict()
    for epoch in range(startEpoch, DefaultConfig.EpochNum):
        since = time.time()
        print('Epoch {}/{}'.format(epoch, DefaultConfig.EpochNum - 1))
        print('-' * 10)

        # Train phase
        model.train()
        running_loss = 0.0
        running_corrects = 0.0
        class_corrects = np.zeros(trashSet.class_num)
        group_corrects = dict()
        for k, v in trashSet.group_index.items():
            group_corrects[k] = 0

        for  i, data in enumerate(dataloaders['train']):
            inputs, labels = data
            inputs = inputs.to(device)
            labels = labels.to(device)

            # zero the parameter gradients
            optimizer.zero_grad()

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels)
            for k, index_list in trashSet.group_index.items():
                group_labels = np.array([1 if x in index_list else 2 for x in labels])
                group_preds = np.array([1 if x in index_list else 3 for x in preds])
                group_corrects[k] += np.sum(group_preds==group_labels)

            for j in range(trashSet.class_num):
                class_labels = np.array([1 if x==j else 2 for x in labels])
                class_preds = np.array([1 if x==j else 3 for x in preds])
                class_corrects[j] += np.sum(class_preds==class_labels)

            print("epoch:", epoch, '/',DefaultConfig.EpochNum ,
                  'step loss', loss.item(),
                  'global_step=', epoch * len(dataloaders['train']) + i)

            if i % 100 == 99:
                writer.add_scalar('train/mini_batch_loss',
                                  loss.item(),
                                  global_step = epoch * len(dataloaders['train']) + i)


        exp_lr_scheduler.step()
        epoch_loss = running_loss / trashSet.dataset_sizes['train']
        epoch_acc = running_corrects.double() / trashSet.dataset_sizes['train']
        class_acc = class_corrects / trashSet.class_train_image_count
        for k, v in group_corrects.items():
            group_acc_dict[k] = group_corrects[k]/trashSet.group_count['train'][k]

        for j in range(trashSet.class_num):
            class_acc_dict[trashSet.index_to_class[j]] = class_acc[j]

        print('Train Loss: {:.4f} Acc: {:.4f}'.format(epoch_loss, epoch_acc))
        writer.add_scalar('train/Epoch loss', epoch_loss, global_step=epoch)
        writer.add_scalar('train/Train Acc', epoch_acc, global_step=epoch)
        writer.add_scalars('train/Class Acc', class_acc_dict, global_step=epoch)
        writer.add_scalars('train/Group Acc', group_acc_dict, global_step=epoch)

        torch.save({
            'epoch': epoch,
            'state_dict': model.state_dict(),
            'epoch_loss': epoch_loss,
            'train_acc': epoch_acc,
            'optimizer_state_dict': optimizer.state_dict()
        }, DefaultConfig.LogDir+'/ckpt_'+str(epoch)+'.pth')
        print('Epoch',epoch, 'model saved.')

        # Validation phase
        model.eval()
        valid_loss = 0.0
        valid_corrects = 0.0
        class_corrects = np.zeros(trashSet.class_num)
        for k, v in trashSet.group_index.items():
            group_corrects[k] = 0

        with torch.no_grad():
            for inputs, labels in dataloaders['val']:
                inputs = inputs.to(device)
                labels = labels.to(device)
                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                loss = criterion(outputs, labels)

                valid_loss += loss.item() * inputs.size(0)
                valid_corrects += torch.sum(preds == labels)

                for k, index_list in trashSet.group_index.items():
                    group_labels = np.array([1 if x in index_list else 2 for x in labels])
                    group_preds = np.array([1 if x in index_list else 3 for x in preds])
                    group_corrects[k] += np.sum(group_preds == group_labels)

                for j in range(trashSet.class_num):
                    class_labels = np.array([1 if x == j else 2 for x in labels])
                    class_preds = np.array([1 if x == j else 3 for x in preds])
                    class_corrects[j] += np.sum(class_preds == class_labels)

            epoch_valid_loss = valid_loss / trashSet.dataset_sizes['val']
            epoch_valid_acc = valid_corrects / trashSet.dataset_sizes['val']
            class_acc = class_corrects / trashSet.class_val_image_count
            for j in range(trashSet.class_num):
                class_acc_dict[trashSet.index_to_class[j]] = class_acc[j]

            for k, v in group_corrects.items():
                group_acc_dict[k] = group_corrects[k] / trashSet.group_count['val'][k]

            print('Valid Loss: {:.4f} Acc: {:.4f}'.format(epoch_valid_loss, epoch_valid_acc))
            writer.add_scalar('valid/Epoch valid loss', epoch_valid_loss, global_step=epoch)
            writer.add_scalar('valid/Valid Acc', epoch_valid_acc, global_step=epoch)
            writer.add_scalars('valid/Class Acc', class_acc_dict, global_step=epoch)
            writer.add_scalars('valid/Group Acc', group_acc_dict, global_step=epoch)

        print('Epoch time cost:', time.time()-since)
        print()
    writer.close()
if __name__ == '__main__':
    trash_set = TrashSet()
    trash_set.save_params(DefaultConfig.DataSetInfoPath)
    train_with_valid(trash_set)