0
0

基于ResNet50实现垃圾分类

一、垃圾分类背景

在现代社会中，垃圾分类已经成为环保的重要措施之一。然而，手动分类垃圾既费时又容易出错。借助深度学习技术，特别是卷积神经网络（CNN），我们可以开发一种自动垃圾分类系统。本文将介绍如何基于ResNet50实现垃圾分类。

二、ResNet50简介

ResNet50是Residual Networks（残差网络）的一种变体，由Kaiming He等人在2015年提出。ResNet50包含50个深度层，通过引入残差模块，有效地解决了深层网络的退化问题。残差模块通过引入短连接（skip connections）使得网络在训练时更容易优化。下图在下文中称为主图

1、ResNet50整体架构图

首先需要声明，这张图的内容是ResNet的Backbone部分（即图中没有ResNet中的全局平均池化层和全连接层），输入

  1. INPUT

经过ResNet50的5个阶段（Stage 0、Stage 1、……）得到输出

  1. OUTPUT

。

下面附上ResNet原文展示的ResNet结构，大家可以结合着看，看不懂也没关系，只看本文也可以无痛理解的。

上图描述了ResNet多个版本的具体结构，本文描述的“ResNet50”中的50指有50个层。和上图一样，本图描述的ResNet也分为5个阶段。

2、ResNet50各个部分的具体架构

1）Stage结构

  • (3,224,224)指输入INPUT的通道数(channel)、高(height)和宽(width)，即(C,H,W)。现假设输入的高度和宽度相等，所以用(C,W,W)表示。
  • 该stage中第1层包括3个先后操作1. CONV``````CONV是卷积（Convolution）的缩写，7×7指卷积核大小，64指卷积核的数量（即该卷积层输出的通道数），/2指卷积核的步长为2。2. BN``````BN是Batch Normalization的缩写，即常说的BN层。3. RELU``````RELU指ReLU激活函数。
  • 该stage中第2层为MAXPOOL，即最大池化层，其kernel大小为3×3、步长为2。
  • (64,56,56)是该stage输出的通道数(channel)、高(height)和宽(width)，其中64等于该stage第1层卷积层中卷积核的数量，56等于224/2/2（步长为2会使输入尺寸减半）。

总体来讲，在Stage 0中，形状为

  1. (3,224,224)

的输入先后经过卷积层、BN层、ReLU激活函数、MaxPooling层得到了形状为

  1. (64,56,56)

的输出。

2）BINK1、BINK2的结构

BINK2（主图右侧部分）：

BINK2有两个参数：C,W

C：代表输入通道数。W：代表输入尺寸。

BINK2左侧经过三个卷积快（包括BN,RELU），设其输出为F（x），将F和x相加再经过Relu激活函数得到BINK2的输出。至于为什么将F和x相加后再输出后面将会介绍。

BINK1：

BINK1有四个参数：C,W,C1,S。

S：代表卷积层中的步长，当S为1时，输入尺寸和输出尺寸相同，代表没有进行下采样。

C1：代表卷积层输出的特征图数目，即输出通道数。

C：代表输入通道数。C和C1相等说明左侧1×1的卷积层没有减少通道数，后三个stage中C=2*C1说明左侧1×1的卷积层减少了通道数。

W：代表输入尺寸，即长和宽。

BINK1相对于BINK2是输入通道和输出通道不一致的情况，BINK1右侧先经过一个卷积层，改变其输出通道数，设其输出为G（x），G函数起到了和左侧输出通道数匹配的作用，这样将F和G相加再经过Relu激活函数得到BINK1的输出。

3）简要分析

原文可知，ResNet后4个stage中都有

  1. BTNK1

和

  1. BTNK2

。

  • 4个stage中BTNK2参数规律相同4个stage中BTNK2的参数全都是1个模式和规律，只是输入的形状(C,W,W)不同。
  • Stage 1中BTNK1参数的规律与后3个stage不同然而，4个stage中BTNK1的参数的模式并非全都一样。具体来讲，后3个stage中BTNK1的参数模式一致，Stage 1中BTNK1的模式与后3个stage的不一样，这表现在以下2个方面：1. 参数S：BTNK1左右两个1×1卷积层是否下采样Stage 1中的BTNK1：步长S为1，没有进行下采样，输入尺寸和输出尺寸相等。后3个stage的BTNK1：步长S为2，进行了下采样，输入尺寸是输出尺寸的2倍。2. 参数C和C1：BTNK1左侧第一个1×1卷积层是否减少通道数Stage 1中的BTNK1：输入通道数C和左侧1×1卷积层通道数C1相等（C=C1=64），即左侧1×1卷积层没有减少通道数。后3个stage的BTNK1：输入通道数C和左侧1×1卷积层通道数C1不相等（C=2*C1），左侧1×1卷积层有减少通道数。
  • 为什么Stage 1中BTNK1参数的规律与后3个stage不同？（个人观点）- 关于BTNK1左右两个1×1卷积层是否下采样因为Stage 0中刚刚对网络输入进行了卷积和最大池化，还没有进行残差学习，此时直接下采样会损失大量信息；而后3个stage直接进行下采样时，前面的网络已经进行过残差学习了，所以可以直接进行下采样。- 关于BTNK1左侧第一个1×1卷积层是否减少通道数根据ResNet原文可知，Bottleneck左侧两个1×1卷积层的主要作用分别是减少通道数和恢复通道数，这样就可以使它们中间的3×3卷积层的输入和输出的通道数都较小，因此效率更高。Stage 1中BTNK1的输入通道数C为64，它本来就比较小，因此没有必要通过左侧第一个1×1卷积层减少通道数。

4）残差结构

传统的卷积神经网络（CNN）在训练过程中，当网络深度增加时，梯度消失和退化问题变得更加明显。为了解决这些问题，残差结构引入了短连接，通过直接将输入跳跃连接到输出，形成了所谓的“残差”连接。具体来说，残差结构通过以下公式来表达：

y=F(x,{Wi})+x\mathbf{y} = \mathcal{F}(\mathbf{x}, {W_i}) + \mathbf{x}y=F(x,{Wi​})+x

其中：

  • x\mathbf{x}x 是输入。
  • F(x,{Wi})\mathcal{F}(\mathbf{x}, {W_i})F(x,{Wi​}) 表示学习到的残差函数，即输入经过若干卷积层后的输出。
  • y\mathbf{y}y 是残差结构的最终输出。
  • {Wi}{W_i}{Wi​} 表示卷积层的权重。

通过这种形式，残差结构能够确保输入信息在每一层都能得到保留和传递，从而缓解了深层网络中的梯度消失问题。

三、ResNet50实现垃圾分类

1、数据集准备

1）首先，我们需要一组垃圾分类的数据集。

常用的数据集有：

  • Garbage Classification Dataset
  • TrashNet

我们以垃圾分类数据集为例，该数据集包含多个类别的垃圾图像，如纸张、塑料、金属等。

本文采用自定义的158类别数据集。

2）数据集的处理

对于文章的数据，我们定义一个split_data.py文件用于生成训练集，验证集和测试集。

  1. import os
  2. from shutil import copy, rmtree
  3. import random
  5. def mk_file(file_path: str):
  6.     if os.path.exists(file_path):
  7.         # 如果文件夹存在，则先删除原文件夹在重新创建
  8.         rmtree(file_path)
  9.     os.makedirs(file_path)
  11. def main():
  12.     # 保证随机可复现
  13.     random.seed(0)
  15.     # 将数据集中10%的数据划分到验证集中
  16.     split_rate = 0.1
  18.     # 将数据集中20%的数据划分到测试集中
  19.     split_rate_2 = 0.2
  21.     #因此训练集、验证集、测试集比例为 7：1：2
  23.     # 指向你解压后的garbage_photos文件夹
  24.     cwd = os.getcwd()
  25.     data_root_pro = os.path.join(cwd, "data_set")  
  26.     data_root = os.path.join(data_root_pro, "garbage_data") # data_root = /data_set/garbage
  27.     assert os.path.exists(data_root), "path '{}' does not exist.".format(data_root)
  29.     # flower_class = [cla for cla in os.listdir(data_root)
  30.     #                 if os.path.isdir(os.path.join(data_root, cla))]
  31.     
  32.     garbage_class = [cla for cla in os.listdir(data_root)
  33.                     if os.path.isdir(os.path.join(data_root, cla))]   
  35.     # 建立保存训练集的文件夹
  36.     train_root = os.path.join(data_root, "train_garbage")
  37.     mk_file(train_root)
  38.     for cla in garbage_class:
  39.         # 建立每个类别对应的文件夹
  40.         mk_file(os.path.join(train_root, cla))
  42.     # 建立保存验证集的文件夹
  43.     val_root = os.path.join(data_root, "val_garbage")
  44.     mk_file(val_root)
  45.     for cla in garbage_class:
  46.         # 建立每个类别对应的文件夹
  47.         mk_file(os.path.join(val_root, cla))
  49.     # 建立保存测试集的文件夹
  50.     test_root = os.path.join(data_root, "test_garbage")  #对于测试集无需进行分类保存图片
  51.     mk_file(test_root)
  53.     
  55.     for cla in garbage_class:
  56.         cla_path = os.path.join(data_root, cla)  # /data_set/garbage/0/
  57.         images = os.listdir(cla_path)
  58.         num = len(images)
  59.         # 随机采样验证集和测试集的索引
  60.         eval_index = random.sample(images, k=int(num*split_rate))
  61.         test_index = random.sample(images, k=int(num*split_rate_2))
  62.         for index, image in enumerate(images):
  63.             if image in eval_index:
  64.                 # 将分配至验证集中的文件复制到相应目录
  65.                 image_path = os.path.join(cla_path, image)
  66.                 new_path = os.path.join(val_root, cla)
  67.                 copy(image_path, new_path)
  68.             elif image in test_index:
  69.                 # 将分配至测试集中的文件复制到相应目录
  70.                 image_path = os.path.join(cla_path, image)
  71.                 #new_path = os.path.join(test_root, cla)  
  72.                 copy(image_path, test_root)
  73.             else:
  74.                 # 将分配至训练集中的文件复制到相应目录
  75.                 image_path = os.path.join(cla_path, image)
  76.                 new_path = os.path.join(train_root, cla)
  77.                 copy(image_path, new_path)
  78.             print("\r[{}] processing [{}/{}]".format(cla, index+1, num), end="")  # processing bar
  79.         print()
  81.     print("processing done!")
  83. if __name__ == '__main__':
  84.     main()

运行split_data.py文件之后，目录将生成下图所示三个文件夹，其中train_garbage和val_garbage包含158个子文件夹，每个文件夹存在对应的图像。

2、 模型的构建

  1. import torch.nn as nn
  2. import torch
  4. class BasicBlock(nn.Module):
  5.     expansion = 1
  6.     # 适用于ResNet18和ResNet34的基本残差块
  7.     def __init__(self, in_channel, out_channel, stride=1, downsample=None, **kwargs):
  8.         super(BasicBlock, self).__init__()
  9.         # 第一个3x3卷积层
  10.         self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
  11.                                kernel_size=3, stride=stride, padding=1, bias=False)
  12.         self.bn1 = nn.BatchNorm2d(out_channel)  # 批量归一化层
  13.         self.relu = nn.ReLU()  # 激活函数
  14.         # 第二个3x3卷积层
  15.         self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
  16.                                kernel_size=3, stride=1, padding=1, bias=False)
  17.         self.bn2 = nn.BatchNorm2d(out_channel)  # 批量归一化层
  18.         self.downsample = downsample  # 下采样层（用于调整输入尺寸和通道数）
  20.     def forward(self, x):
  21.         identity = x  # 保存输入值
  22.         if self.downsample is not None:  # 如果需要下采样
  23.             identity = self.downsample(x)
  25.         out = self.conv1(x)  # 第一个卷积层
  26.         out = self.bn1(out)  # 批量归一化
  27.         out = self.relu(out)  # 激活
  29.         out = self.conv2(out)  # 第二个卷积层
  30.         out = self.bn2(out)  # 批量归一化
  32.         out += identity  # 残差连接
  33.         out = self.relu(out)  # 激活
  35.         return out
  37. class Bottleneck(nn.Module):
  38.     expansion = 4
  39.     # 适用于ResNet50、ResNet101和ResNet152的瓶颈残差块
  40.     def __init__(self, in_channel, out_channel, stride=1, downsample=None,
  41.                  groups=1, width_per_group=64):
  42.         super(Bottleneck, self).__init__()
  44.         width = int(out_channel * (width_per_group / 64.)) * groups  # 计算组卷积的宽度
  46.         # 第一个1x1卷积层（压缩通道数）
  47.         self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=width,
  48.                                kernel_size=1, stride=1, bias=False)
  49.         self.bn1 = nn.BatchNorm2d(width)  # 批量归一化层
  50.         # 第二个3x3卷积层
  51.         self.conv2 = nn.Conv2d(in_channels=width, out_channels=width, groups=groups,
  52.                                kernel_size=3, stride=stride, bias=False, padding=1)
  53.         self.bn2 = nn.BatchNorm2d(width)  # 批量归一化层
  54.         # 第三个1x1卷积层（扩展通道数）
  55.         self.conv3 = nn.Conv2d(in_channels=width, out_channels=out_channel*self.expansion,
  56.                                kernel_size=1, stride=1, bias=False)
  57.         self.bn3 = nn.BatchNorm2d(out_channel*self.expansion)  # 批量归一化层
  59.         self.relu = nn.ReLU(inplace=True)  # 激活函数
  60.         self.downsample = downsample  # 下采样层
  62.     def forward(self, x):
  63.         identity = x  # 保存输入值
  64.         if self.downsample is not None:  # 如果需要下采样
  65.             identity = self.downsample(x)
  67.         out = self.conv1(x)  # 第一个卷积层
  68.         out = self.bn1(out)  # 批量归一化
  69.         out = self.relu(out)  # 激活
  71.         out = self.conv2(out)  # 第二个卷积层
  72.         out = self.bn2(out)  # 批量归一化
  73.         out = self.relu(out)  # 激活
  75.         out = self.conv3(out)  # 第三个卷积层
  76.         out = self.bn3(out)  # 批量归一化
  78.         out += identity  # 残差连接
  79.         out = self.relu(out)  # 激活
  81.         return out
  83. class ResNet(nn.Module):
  85.     def __init__(self,
  86.                  block,
  87.                  blocks_num,  # 每个阶段的残差块数量
  88.                  num_classes=1000,
  89.                  include_top=True,
  90.                  groups=1,
  91.                  width_per_group=64):
  92.         super(ResNet, self).__init__()
  93.         self.include_top = include_top  # 是否包含全连接层
  94.         self.in_channel = 64  # 初始通道数
  96.         self.groups = groups  # 组卷积数量
  97.         self.width_per_group = width_per_group  # 每组的宽度
  99.         # 初始卷积层
  100.         self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2,
  101.                                padding=3, bias=False)
  102.         self.bn1 = nn.BatchNorm2d(self.in_channel)  # 批量归一化层
  103.         self.relu = nn.ReLU(inplace=True)  # 激活函数
  104.         # 最大池化层
  105.         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
  107.         # 各个阶段的残差块
  108.         self.layer1 = self._make_layer(block, 64, blocks_num[0])
  109.         self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)
  110.         self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)
  111.         self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)
  112.         
  113.         # 全局平均池化层和全连接层
  114.         if self.include_top:
  115.             self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # 自适应平均池化层
  116.             self.fc = nn.Linear(512 * block.expansion, num_classes)  # 全连接层
  118.         # 初始化卷积层权重
  119.         for m in self.modules():
  120.             if isinstance(m, nn.Conv2d):
  121.                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
  123.     def _make_layer(self, block, channel, block_num, stride=1):
  124.         downsample = None  # 对于ResNet18和ResNet34，下采样默认为None，其他层为下面函数
  125.         if stride != 1 or self.in_channel != channel * block.expansion:  # 如果步长不为1或通道数不匹配
  126.             downsample = nn.Sequential(
  127.                 nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
  128.                 nn.BatchNorm2d(channel * block.expansion))  # 下采样层
  130.         layers = []
  131.         layers.append(block(self.in_channel,
  132.                             channel,  # 主分支上第一个卷积核（通道）的个数
  133.                             downsample=downsample,  # 下采样层
  134.                             stride=stride,
  135.                             groups=self.groups,
  136.                             width_per_group=self.width_per_group))
  137.         self.in_channel = channel * block.expansion  # 更新输入通道数
  139.         for _ in range(1, block_num):  # 添加剩余的残差块
  140.             layers.append(block(self.in_channel,
  141.                                 channel,
  142.                                 groups=self.groups,
  143.                                 width_per_group=self.width_per_group))
  145.         return nn.Sequential(*layers)  # 将layers列表转化为nn.Sequential
  147.     def forward(self, x):
  148.         x = self.conv1(x)  # 初始卷积层
  149.         x = self.bn1(x)  # 批量归一化
  150.         x = self.relu(x)  # 激活
  151.         x = self.maxpool(x)  # 最大池化层
  153.         x = self.layer1(x)  # 第一阶段
  154.         x = self.layer2(x)  # 第二阶段
  155.         x = self.layer3(x)  # 第三阶段
  156.         x = self.layer4(x)  # 第四阶段
  158.         if self.include_top:  # 如果包含全连接层
  159.             x = self.avgpool(x)  # 全局平均池化层
  160.             x = torch.flatten(x, 1)  # 展平
  161.             x = self.fc(x)  # 全连接层
  163.         return x
  165. def resnet50(model_name="resnet50", num_classes=1000, init_weights=False, **kwargs):
  166.     # 创建ResNet-50模型
  167.     model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, include_top=True)
  168.     
  169.     # 替换最后的全连接层以适应新的类别数
  170.     model.fc = nn.Linear(model.fc.in_features, num_classes)
  171.     
  172.     # 初始化权重（如果需要）
  173.     if init_weights:
  174.         def init_weights(m):
  175.             if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
  176.                 nn.init.kaiming_normal_(m.weight)
  177.                 if m.bias is not None:
  178.                     nn.init.constant_(m.bias,

3、训练模型

  1. # -*- coding: utf-8 -*-
  2. import os
  3. import sys
  4. import json
  5. import torch
  6. import torch.nn as nn
  7. import torch.optim as optim
  8. from torchvision import transforms, datasets
  9. from tqdm import tqdm
  11. from model import resnet50  # 导入定义的ResNet50模型
  13. def main():
  14.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 设置运行设备
  15.     print("using {} device.".format(device))
  17.     # 数据预处理
  18.     data_transform = {
  19.         "train_garbage": transforms.Compose([transforms.RandomResizedCrop(224),  # 随机裁剪并调整大小
  20.                                              transforms.RandomHorizontalFlip(),  # 随机水平翻转
  21.                                              transforms.ToTensor(),  # 转换为张量
  22.                                              transforms.Normalize([0.485, 0.456, 0.406], 
  23.                                                                   [0.229, 0.224, 0.225])]),  # 归一化
  24.         "val_garbage": transforms.Compose([transforms.Resize(256),  # 调整大小
  25.                                            transforms.CenterCrop(224),  # 中心裁剪
  26.                                            transforms.ToTensor(),  # 转换为张量
  27.                                            transforms.Normalize([0.485, 0.456, 0.406], 
  28.                                                                 [0.229, 0.224, 0.225])])}  # 归一化
  30.     data_root = os.path.abspath(os.path.join(os.getcwd(), "./"))  # 获取数据根路径
  31.     print(data_root)
  32.     image_path = os.path.join(data_root, "data_set", "garbage_data")  # 垃圾数据集路径
  33.     assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
  34.     train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train_garbage"),
  35.                                          transform=data_transform["train_garbage"])  # 加载训练数据集
  36.     train_num = len(train_dataset)
  38.     # 获取类别映射字典并写入json文件
  39.     garbage_list = train_dataset.class_to_idx
  40.     cla_dict = dict((val, key) for key, val in garbage_list.items())
  41.     json_str = json.dumps(cla_dict, indent=4)
  42.     with open('/home/dell/CV408/hb/data_set/garbage_data/garbage_classification.json', 'w') as json_file:
  43.         json_file.write(json_str)
  45.     batch_size = 32
  46.     nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # 设置dataloader的worker数量
  47.     print('Using {} dataloader workers every process'.format(nw))
  49.     train_loader = torch.utils.data.DataLoader(train_dataset,
  50.                                                batch_size=batch_size, shuffle=True,
  51.                                                num_workers=nw)  # 加载训练数据集
  53.     validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val_garbage"),
  54.                                             transform=data_transform["val_garbage"])  # 加载验证数据集
  55.     val_num = len(validate_dataset)
  56.     validate_loader = torch.utils.data.DataLoader(validate_dataset,
  57.                                                   batch_size=batch_size, shuffle=False,
  58.                                                   num_workers=nw)  # 加载验证数据集
  60.     print("using {} images for training, {} images for validation.".format(train_num,
  61.                                                                            val_num))
  62.     
  63.     # 加载ResNet50模型
  64.     model_name = "resnet50"
  65.     net = resnet50(model_name=model_name, num_classes=158, init_weights=True)
  66.     net.to(device)  # 将模型移至指定设备
  68.     # 定义损失函数
  69.     loss_function = nn.CrossEntropyLoss()
  71.     # 定义优化器
  72.     params = [p for p in net.parameters() if p.requires_grad]
  73.     optimizer = optim.Adam(params, lr=0.0001)
  75.     epochs = 30
  76.     best_acc = 0.0 
  77.     save_path = './Test5_resnet{}.pth'.format(model_name)  # 保存当前最好的权重路径
  78.     train_steps = len(train_loader)
  80.     for epoch in range(epochs):
  81.         # 训练过程
  82.         net.train()
  83.         running_loss = 0.0
  84.         train_bar = tqdm(train_loader, file=sys.stdout)
  85.         for step, data in enumerate(train_bar):
  86.             images, labels = data
  87.             optimizer.zero_grad()
  88.             logits = net(images.to(device))
  89.             loss = loss_function(logits, labels.to(device))
  90.             loss.backward()
  91.             optimizer.step()
  93.             # 打印损失
  94.             running_loss += loss.item()
  95.             train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
  96.                                                                      epochs,
  97.                                                                      loss)
  99.         # 验证过程
  100.         net.eval()
  101.         acc = 0.0  # 累积准确数
  102.         with torch.no_grad():
  103.             val_bar = tqdm(validate_loader, file=sys.stdout)
  104.             for val_data in val_bar:
  105.                 val_images, val_labels = val_data
  106.                 outputs = net(val_images.to(device))
  107.                 predict_y = torch.max(outputs, dim=1)[1]
  108.                 acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
  110.                 val_bar.desc = "valid epoch[{}/{}]".format(epoch + 1,
  111.                                                            epochs)
  113.         val_accurate = acc / val_num
  114.         print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
  115.               (epoch + 1, running_loss / train_steps, val_accurate))
  117.         if val_accurate > best_acc:  # 更新最好权重
  118.             best_acc = val_accurate
  119.             torch.save(net.state_dict(), save_path)
  121.     print('Finished Training')
  123. if __name__ == '__main__':
  124.     main()

对模型进行预训练之后，将得到一个最好的权重文件 。

4、测试模型

1）单图片测试

  1. import os
  2. import json
  4. import torch
  5. from PIL import Image
  6. from torchvision import transforms
  7. import matplotlib.pyplot as plt
  9. from model import resnet50  # 导入定义的ResNet50模型
  11. def main():
  12.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 设置运行设备
  14.     # 定义数据预处理步骤
  15.     data_transform = transforms.Compose(
  16.         [transforms.Resize(256),  # 调整大小到256x256
  17.          transforms.CenterCrop(224),  # 中心裁剪到224x224
  18.          transforms.ToTensor(),  # 转换为张量
  19.          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])  # 归一化
  21.     # 加载图像
  22.     img_path = "/home/dell/CV408/hb/data_set/garbage_data/test_garbage/8/1171.jpg"  # 图像路径
  23.     assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
  24.     img = Image.open(img_path)  # 打开图像
  25.     plt.imshow(img)  # 显示图像
  26.     img = data_transform(img)  # 对图像进行预处理
  27.     img = torch.unsqueeze(img, dim=0)  # 扩展批次维度
  29.     # 读取类别字典
  30.     json_path = '/home/dell/CV408/hb/data_set/garbage_data/garbage_classification.json'
  31.     assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
  32.     with open(json_path, "r") as f:
  33.         class_indict = json.load(f)
  35.     # 创建模型
  36.     model = resnet50(num_classes=158).to(device)  # 初始化ResNet50模型并设置类别数为158
  38.     # 加载模型权重
  39.     weights_path = "/home/dell/CV408/hb/resnet50.pth"  # 模型权重路径
  40.     assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
  41.     model.load_state_dict(torch.load(weights_path, map_location=device))  # 加载权重
  43.     # 创建保存预测结果的目录
  44.     save_dir = "/home/dell/CV408/hb/data_set/garbage_data/"
  45.     os.makedirs(save_dir, exist_ok=True)
  47.     # 预测过程
  48.     model.eval()  # 设置模型为评估模式
  49.     with torch.no_grad():  # 禁用梯度计算
  50.         output = torch.squeeze(model(img.to(device))).cpu()  # 模型预测并去除批次维度
  51.         predict = torch.softmax(output, dim=0)  # 应用softmax获取概率
  52.         predict_cla = torch.argmax(predict).numpy()  # 获取预测类别
  54.     # 打印预测结果
  55.     print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
  56.                                                  predict[predict_cla].numpy())
  57.     plt.title(print_res)  # 在图像上显示预测结果
  58.     for i in range(len(predict)):  # 打印所有类别的概率
  59.         print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
  60.                                                   predict[i].numpy()))
  61.         
  62.     plt.savefig(os.path.join(save_dir, "predicted_image_garbage.png"))  # 保存带有预测结果的图像
  63.     plt.show()  # 显示图像
  65. if __name__ == '__main__':
  66.     main()  # 运行main函数

测试结果如下图（类别144为：火龙果）

2）批量测试

  1. import os
  2. import json
  3. import matplotlib.pyplot as plt
  4. import torch
  5. from PIL import Image           
  6. from torchvision import transforms
  8. from model import resnet50
  10. def main():
  11.     # 设置设备为GPU（如果可用）或CPU
  12.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  14.     # 定义数据预处理操作，包括调整大小、中心裁剪、转换为张量和归一化
  15.     data_transform = transforms.Compose([
  16.         transforms.Resize(256),
  17.         transforms.CenterCrop(224),
  18.         transforms.ToTensor(),
  19.         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  20.     ])
  22.     # 指定需要遍历预测的图像文件夹
  23.     imgs_root = "/home/dell/CV408/hb/data_set/garbage_data/test_garbage"
  24.     assert os.path.exists(imgs_root), f"file: '{imgs_root}' does not exist."
  26.     # 读取指定文件夹下所有jpg图像路径
  27.     img_path_list = [os.path.join(imgs_root, i) for i in os.listdir(imgs_root) if i.endswith(".jpg")]
  29.     # 读取类别映射文件
  30.     json_path = '/home/dell/CV408/hb/data_set/garbage_data/garbage_classification.json'
  31.     assert os.path.exists(json_path), f"file: '{json_path}' does not exist."
  33.     with open(json_path, "r") as json_file:
  34.         class_indict = json.load(json_file)
  36.     # 创建ResNet50模型并加载预训练权重
  37.     model = resnet50(num_classes=158).to(device)
  38.     weights_path = "/home/dell/CV408/hb/Test5_resnetresnet50.pth"
  39.     assert os.path.exists(weights_path), f"file: '{weights_path}' does not exist."
  40.     model.load_state_dict(torch.load(weights_path, map_location=device))
  42.     # 设置模型为评估模式
  43.     model.eval()
  45.     batch_size = 8  # 每次预测时处理的图像数量
  46.     save_dir = "/home/dell/CV408/hb/data_set/garbage_data/predictions/"
  47.     os.makedirs(save_dir, exist_ok=True)  # 创建保存预测结果的目录
  49.     with torch.no_grad():
  50.         # 按批次处理图像
  51.         for ids in range(0, len(img_path_list) // batch_size):
  52.             img_list = []
  53.             img_paths = img_path_list[ids * batch_size: (ids + 1) * batch_size]
  54.             for img_path in img_paths:
  55.                 assert os.path.exists(img_path), f"file: '{img_path}' does not exist."
  56.                 img = Image.open(img_path)
  57.                 img = data_transform(img)
  58.                 img_list.append(img)
  60.             # 将图像列表打包成一个批次
  61.             batch_img = torch.stack(img_list, dim=0)
  62.             # 预测类别
  63.             output = model(batch_img.to(device)).cpu()
  64.             predict = torch.softmax(output, dim=1)
  65.             probs, classes = torch.max(predict, dim=1)
  67.             # 可视化并保存每张图像的预测结果
  68.             for idx, (pro, cla) in enumerate(zip(probs, classes)):
  69.                 img_path = img_paths[idx]
  70.                 img = Image.open(img_path)
  71.                 plt.imshow(img)
  72.                 plt.title(f"Class: {class_indict[str(cla.numpy())]}  Prob: {pro.numpy():.3f}")
  73.                 save_path = os.path.join(save_dir, f"pred_{os.path.basename(img_path)}")
  74.                 plt.savefig(save_path)
  75.                 plt.close()
  76.                 print(f"image: {img_path}  class: {class_indict[str(cla.numpy())]}  prob: {pro.numpy():.3f}")
  78. if __name__ == '__main__':
  79.     main()

生成一个文件夹保存着测试集图像对应的预测结果。

五、总结

模型的精确率达到了69.7%左右，并且对一些样本少的类别和图像质量较差的预测结果不理想，可能需要对图像的预处理方法进行改进，比如图像增强等等

标签： 分类 数据挖掘 人工智能
本文转载自: https://blog.csdn.net/qq_52191127/article/details/140814422
版权归原作者 只爱喝水 所有， 如有侵权，请联系我们删除。
登录后发布评论

“基于ResNet50实现垃圾分类”的评论:

还没有评论
关于作者
...
overfit同步小助手
文章同步
相关阅读

网络安全法-网络运行安全

使用selenium/drissionpage时如何阻止chrome自动跳转http到https

docker desktop 里部署的Open WebUI 管理员密码忘记了的处理方法

在ubuntu20.04中搭建onsite比赛运行环境

利用开源的低代码表单设计器FcDesigner高效管理和渲染复杂表单结构

Kafka学习笔记

【前端】浏览器输入url到页面呈现发生了什么？