CS课程项目设计20：手把手教你实现垃圾分类系统

最新推荐文章于 2025-11-14 02:47:04 发布

原创最新推荐文章于 2025-11-14 02:47:04 发布 · 1.8k 阅读

36 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #深度学习 #Python #图像处理 #垃圾分类系统

CS课程项目设计专栏收录该内容

28 篇文章

订阅专栏

1. 研究背景

随着城市化进程加快和居民生活水平提高，生活垃圾产量急剧增加，垃圾分类成为提升资源回收利用率、减少环境污染的关键环节。传统垃圾分类依赖人工识别，存在效率低、主观性强、分类标准不统一等问题。

近年来，深度学习技术在图像识别领域取得突破性进展，为垃圾分类自动化提供了可能。ResNet（残差网络）作为深度卷积神经网络的典型代表，凭借其解决深层网络梯度消失问题的独特设计，在图像分类任务中表现优异，成为垃圾分类系统的理想技术选择。

2. 研究目的

本系统旨在通过 ResNet 图像识别模型实现生活垃圾的自动分类，具体目标包括：

实现六大类垃圾的精准识别： cardboard（纸板）、glass（玻璃）、metal（金属）、paper（纸张）、plastic（塑料）、other_trash（其他垃圾）
提供高效的单张图像识别功能，支持常见图像格式（jpg、jpeg、png、bmp）的快速分类
支持批量图像处理，满足大量垃圾图像的批量分类需求，并生成结构化结果报告
构建用户友好的可视化界面，降低垃圾分类技术的使用门槛
实现模型的自主训练功能，支持用户根据实际需求优化模型性能
达到较高的分类准确率（目标验证集准确率 > 85%），满足实际应用场景需求

六大类垃圾的图片是从TrashNet垃圾分类数据集中收集，这是一个专门为垃圾分类任务设计的数据集，可通过以下链接访问：

TrashNet垃圾分类数据集链接

六大类垃圾的示例图片如下所示：

3. 技术方案

3.1 核心技术栈

PyTorch 深度学习框架：提供模型构建、训练和推理的完整生态支持，支持 GPU 加速计算，提升模型训练和预测效率。
PIL（Pillow）和torchvision 图像处理库：实现图像加载、格式转换、预处理等功能；提供丰富的图像变换工具（裁剪、旋转、归一化等）。
Tkinter 用户界面框架：构建跨平台的图形用户界面，实现图像预览、结果展示、进度跟踪等交互功能。
Matplotlib 可视化工具：绘制训练过程中的损失和准确率曲线，辅助分析模型训练效果。

3.2 ResNet 核心运行机制

ResNet（残差网络）的核心创新是引入残差块（Residual Block）结构，解决了深层神经网络训练中的梯度消失问题：

一、残差连接设计

每个残差块包含主路径（卷积操作）和跳跃连接（恒等映射或 1×1 卷积调整维度），输出为两者之和：H(x) = F(x) + x，其中 F (x) 为主路径变换，x 为输入的恒等映射。

二、深度网络构建

系统采用 ResNet50 架构，包含 49 个卷积层和 1 个全连接层，通过堆叠残差块实现深度网络结构。

三、迁移学习应用

优先加载本地预训练模型权重，若无则使用 PyTorch 提供的预训练 ResNet50 权重；
冻结大部分网络参数（仅保留最后 10 层可训练），降低训练成本；
替换最后一层全连接层，适配 6 类垃圾的分类任务。

四、特征提取流程

输入图像通过卷积层和池化层逐步提取低级特征（边缘、纹理）到高级特征（形状、语义）；
最终通过全连接层输出分类概率分布，实现垃圾类别的预测。

4. 实现流程

4.1 数据准备阶段

定义数据集目录结构，按类别组织图像数据
实现数据预处理管道，包括训练集增强（随机裁剪、翻转、旋转等）和验证集标准化
按 8:2 比例随机划分训练集和验证集

4.2 模型构建阶段

初始化 ResNet50 基础架构
加载本地预训练模型或官方预训练权重
冻结部分参数，修改输出层适配垃圾分类任务
将模型部署到 GPU（优先）或 CPU 设备

该阶段的核心代码如下所示：

def build_model(num_classes):
    """构建分类模型，自动检测并使用本地ResNet50模型或预训练模型"""
    # 创建基础模型架构
    model = models.resnet50(pretrained=False)  # 先不加载预训练权重

    # 检查本地是否存在ResNet50模型文件
    if os.path.exists(config.local_resnet_path):
        try:
            # 加载本地ResNet50模型权重
            resnet_weights = torch.load(config.local_resnet_path, map_location=config.device)
            model.load_state_dict(resnet_weights)
            print(f"成功加载本地ResNet50模型: {config.local_resnet_path}")
            return_message = f"使用本地ResNet50模型: {config.local_resnet_path}"
        except Exception as e:
            print(f"加载本地ResNet50模型失败: {str(e)}，将使用预训练模型")
            model = models.resnet50(pretrained=True)
            return_message = f"本地模型加载失败，使用预训练ResNet50模型: {str(e)}"
    else:
        # 本地模型不存在，使用预训练的ResNet50权重
        model = models.resnet50(pretrained=True)
        print(f"本地ResNet50模型不存在({config.local_resnet_path})，使用预训练模型")
        return_message = f"本地ResNet50模型不存在，使用预训练模型"

    # 冻结大部分参数
    for param in list(model.parameters())[:-10]:
        param.requires_grad = False

    # 替换最后一层以适应我们的分类任务
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)

    # 移动模型到适当的设备
    model = model.to(config.device)

    return model, return_message

4.3 模型训练阶段

使用交叉熵损失函数计算分类损失
采用 Adam 优化器更新可训练参数
迭代训练10 个 epochs，跟踪训练集和验证集的损失与准确率
保存验证集准确率最高的模型权重

点击“训练模型”按钮后，等进度条走完，训练过程就结束了，如下图所示：

模型训练阶段，只需等待，其核心源代码如下所示：

# 5. 训练和评估函数
def train_model(model, train_loader, val_loader, criterion, optimizer, epochs, progress_callback=None):
    """训练模型"""
    best_acc = 0.0
    train_losses = []
    val_losses = []
    train_accs = []
    val_accs = []

    for epoch in range(epochs):
        print(f'Epoch {epoch + 1}/{epochs}')
        print('-' * 10)

        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0

        for inputs, labels in train_loader:
            inputs = inputs.to(config.device)
            labels = labels.to(config.device)

            # 清零梯度
            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.data)

        epoch_loss = running_loss / len(train_loader.dataset)
        epoch_acc = running_corrects.double() / len(train_loader.dataset)

        train_losses.append(epoch_loss)
        train_accs.append(epoch_acc.item())

        print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

训练结束后，训练损失和准确率指标会保存在本地training_curves.png中，如下图所示：