垃圾分类项目1.0

垃圾分类项目1.0

前言

在深度学习中实现垃圾分类的主要步骤包括数据准备、模型构建、模型训练、模型评估和部署。以下是每个步骤的详细说明：
一、数据准备:
数据收集：收集包含不同类别垃圾的图像数据。
数据清洗与标注：清洗数据以去除无关或重复的图像，并对每个图像进行标注，确保每张图片都分配到正确的垃圾类别。
数据增强：使用数据增强技术（如旋转、裁剪、缩放等）增加数据集的多样性，减少模型过拟合的风险。
数据划分：将数据集划分为训练集、验证集和测试集。
二、模型构建:
选择模型架构：选择合适的神经网络架构（如卷积神经网络CNN、残差网络ResNet等）用于分类任务。
模型定义：使用深度学习框架（如PyTorch、TensorFlow等）定义模型的各层结构和参数。
模型训练:
定义损失函数和优化器：选择适当的损失函数（如交叉熵损失）和优化器（如Adam、SGD等）。
训练循环：编写训练循环代码，在训练集中迭代模型，并根据损失函数最小化的目标调整模型参数。
模型检查点：在训练过程中定期保存模型检查点，以防止训练中断时丢失进度。
三、模型评估:
验证集评估：在每个训练周期后使用验证集评估模型性能，调整超参数以提高模型准确性。
测试集评估：在最终训练结束后，在测试集中评估模型的性能，确保模型在未见过的数据上也能表现良好。

一、数据准备

首先由于刚开始学习，先做一个简单的垃圾分类的小项目，所以跳过数据的清洗和增强，我们对数据进行划分，划为训练集和测试集。

import os
import random

train_ratio = 0.9
rootdata = "D:/Administrator/Desktop/all/all"

train_list, test_list = [], []

class_flag = 0  # 从0开始标记类别

# 遍历根目录
for root, dirs, files in os.walk(rootdata):
    # 遍历每个类别文件夹中的所有文件
    for file in files:
        file_path = os.path.join(root, file)
        # 根据训练和测试比例划分数据
        if random.random() < train_ratio:
            train_list.append(f"{file_path}\t{class_flag}\n")
        else:
            test_list.append(f"{file_path}\t{class_flag}\n")

    # 处理完一个文件夹中的所有文件后，切换到下一个类别
    class_flag += 1

# 打乱列表顺序
random.shuffle(train_list)
random.shuffle(test_list)

# 将训练和测试数据写入文件
with open('train.txt', 'w', encoding='UTF-8') as f:
    f.writelines(train_list)

with open('test.txt', 'w', encoding='UTF-8') as f:
    f.writelines(test_list)

1.变量定义：
train_ratio： 定义训练集占总数据的比例。
rootdata ：是数据集的根目录。
train_list 和 test_list： 分别是存储训练集和测试集的文件路径及其标签。
class_flag： 用于标记当前类别，从0开始。
2.数据划分：
使用 os.walk() 遍历 rootdata 目录及其子目录。
对每个文件，根据 train_ratio 将文件路径和类别标签添加到 train_list 或 test_list。
3.数据打乱：
使用 random.shuffle() 打乱训练集和测试集的顺序，以确保训练过程中数据的随机性。
(作用：防止过拟合，打乱数据可以防止模型过拟合于某个特定的顺序模式。如果数据集中存在某种模式（例如，所有相同类别的图像都集中在一起），模型可能会过度拟合这种模式。通过打乱数据，模型会在训练过程中看到更广泛的组合，从而减少过拟合的风险。）
假设你的数据集中有不同类别的垃圾图片，按照类别存储在文件夹中。如果不打乱数据，模型可能会首先看到所有类别1的图片，然后是类别2的图片，依此类推。这种情况下，模型在初期会集中学习某一类别的特征，而忽略其他类别的特征，导致模型的初期学习效率低下。
通过打乱数据，模型在每个训练周期中都会看到不同类别的混合数据，这样可以帮助模型同时学习不同类别的特征，提高训练效率和模型性能。
4.写入文件：
将训练集和测试集的数据分别写入 train.txt 和 test.txt 文件中

二、模型构建

2.1神经网络构建

在model文件夹下创建net文件，开始搭建自己的网络结构

import torch
import torch.nn as nn
import torch.nn.functional as F

class simpleconv3(nn.Module):
    def __init__(self, num_classes):
        super(simpleconv3, self).__init__()
        self.conv1 = nn.Conv2d(3, 12, 3, 2)  # 输入图片大小为3*48*48，输出特征图大小为12*23*23，卷积核大小为3*3，步长为2
        self.bn1 = nn.BatchNorm2d(12)
        self.conv2 = nn.Conv2d(12, 24, 3, 2)  # 输入图片大小为12*23*23，输出特征图大小为24*11*11，卷积核大小为3*3，步长为2
        self.bn2 = nn.BatchNorm2d(24)
        self.conv3 = nn.Conv2d(24, 48, 3, 2)  # 输入图片大小为24*11*11，输出特征图大小为48*5*5，卷积核大小为3*3，步长为2
        self.bn3 = nn.BatchNorm2d(48)
        self.fc1 = nn.Linear(48 * 5 * 5, 1200)  # 输入向量长为48*5*5=1200，输出向量长为1200
        self.fc2 = nn.Linear(1200, 128)  # 输入向量长为1200，输出向量长为128
        self.fc3 = nn.Linear(128, num_classes)  # 输入向量长为128，输出向量长为num_classes，等于类别数

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = x.view(x.size(0), -1)  # 动态调整批次大小
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

1.定义模型结构：
simpleconv3类继承自nn.Module，在其构造函数__init__中定义了模型的层次结构，包括卷积层、批量归一化层和全连接层。
conv1、conv2、conv3是三层卷积层，分别处理输入图像并提取特征。
bn1、bn2、bn3是对应的批量归一化层，帮助稳定和加速训练过程。
fc1、fc2、fc3是全连接层，负责将卷积特征转换为最终的输出类别。
2.前向传播（forward pass）：
在forward方法中，数据依次通过卷积、批量归一化、ReLU激活函数和全连接层。
x.view(x.size(0), -1)将卷积层的输出展平成一维向量，以便传递给全连接层。

2.2数据加载器

数据加载器的作用
1.批量加载数据：
数据加载器可以将数据集分成多个小批次（batch），在训练和测试过程中逐批加载数据。这对于内存管理和计算效率非常重要。

2.打乱数据：
训练时打乱数据（shuffle）可以防止模型在训练过程中对数据的顺序产生依赖，从而提高模型的泛化能力。

3.并行加载数据：
数据加载器支持多线程并行加载数据，从而加快数据读取速度，减少数据加载时间对训练过程的影响。

4.数据预处理：
数据加载器可以与数据转换操作（如数据增强、标准化等）配合使用，在每次加载数据时对数据进行预处理。

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image, UnidentifiedImageError

class CustomDataset(Dataset):
    def __init__(self, txt_file, transform=None):
        self.img_labels = []
        with open(txt_file, 'r') as file:
            for line in file:
                path, label = line.strip().split('\t')
                # 确保路径格式正确
                path = os.path.normpath(path)
                # 只添加有效的图像路径
                if self._is_image_file(path):
                    self.img_labels.append((path, int(label)))
        self.transform = transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path, label = self.img_labels[idx]
        try:
            image = Image.open(img_path).convert('RGB')  # 确保图像是 RGB 模式
        except (UnidentifiedImageError, IOError) as e:
            print(f"Error loading image {img_path}: {e}")
            # 如果遇到错误，可以选择返回一个空图像或某种默认值
            image = Image.new('RGB', (48, 48))  # 创建一个指定大小的黑色图像
        if self.transform:
            image = self.transform(image)
        return image, label

    def _is_image_file(self, path):
        # 判断文件是否是图像文件
        IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff'}
        return any(path.lower().endswith(ext) for ext in IMAGE_EXTENSIONS)

class TestDataset(Dataset):
    def __init__(self, txt_file, transform=None):
        self.img_paths = []
        with open(txt_file, 'r') as file:
            for line in file:
                path = line.strip()
                # 确保路径格式正确
                path = os.path.normpath(path)
                self.img_paths.append(path)
        self.transform = transform

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        try:
            image = Image.open(img_path).convert('RGB')  # 确保图像是 RGB 模式
        except (UnidentifiedImageError, IOError) as e:
            print(f"Error loading image {img_path}: {e}")
            # 如果遇到错误，可以选择返回一个空图像或某种默认值
            image = Image.new('RGB', (48, 48))  # 创建一个指定大小的黑色图像
        if self.transform:
            image = self.transform(image)
        return image

# 数据转换操作
transform = transforms.Compose([
    transforms.Resize((48, 48)),
    transforms.ToTensor(),
])

# 创建训练集和测试集数据集及数据加载器
train_dataset = CustomDataset(txt_file=r'E:\deep lerning\lajifenlei\train.txt', transform=transform)
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)

test_dataset = CustomDataset(txt_file=r'E:\deep lerning\lajifenlei\test.txt', transform=transform)
test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False)

三、模型训练

创建train.py文件
以下是完整的训练流程：
1.导入必要的模块
2.定义训练参数
3.创建模型实例
4.定义损失函数和优化器
5.训练模型

import torch
from torch import nn
import torch.optim as optim
from model.net import simpleconv3
from dataset.dataset import train_dataloader


def train_model(model, dataloader, num_epochs=5, checkpoint_interval=1, final_model_path="final_model.pth"):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0

        for images, labels in dataloader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        avg_loss = running_loss / len(dataloader)
        print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.4f}")

        # 保存检查点
       # if (epoch + 1) % checkpoint_interval == 0:
            #save_model(model, f"model_checkpoint_epoch_{epoch + 1}.pth")

    # 在训练完成后保存最终模型
    save_model(model, final_model_path)
    print("Training Finished")


def save_model(model, path):
    torch.save(model.state_dict(), path)
    print(f"Model saved to {path}")


def load_model(model, path):
    model.load_state_dict(torch.load(path))
    model.eval()
    print(f"Model weights loaded from {path}")


# 示例用法
if __name__ == "__main__":
    model = simpleconv3(num_classes=50)
    train_model(model, train_dataloader, num_epochs=3, checkpoint_interval=1, final_model_path="final_model.pth")

注意：
from model.net import simpleconv3是将同目录下的model文件夹下的model文件的simpleconv3模型导入，from dataset.dataset import train_dataloadershi从 dataset 模块中的 dataset 文件中导入 train_dataloader 对象。

四、模型评估

from model.net import simpleconv3
from dataset.dataset import test_dataloader
import torch

# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def evaluate_model(model, dataloader):
    model.to(device)  # 确保模型在正确的设备上
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in dataloader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Accuracy of the network on the test images: {100 * correct // total}%')

def load_model(model, file_path):
    model.load_state_dict(torch.load(file_path, map_location=device))
    model.to(device)  # 确保模型在正确的设备上
    model.eval()  # 设置模型为评估模式
    print(f"Model weights loaded from {file_path}")

# 实例化模型
model = simpleconv3(num_classes=50)

# 加载模型权重
load_model(model, 'final_model.pth')

# 评估模型
evaluate_model(model, test_dataloader)

应注意的是model = simpleconv3(num_classes=50)中num_classes应为具体的类数

五、单图片预测

import torch
from PIL import Image
from torchvision import transforms
from model.net import simpleconv3

# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载模型
def load_model(model, file_path):
    model.load_state_dict(torch.load(file_path, map_location=device))
    model.to(device)  # 确保模型在正确的设备上
    model.eval()  # 设置模型为评估模式
    print(f"Model weights loaded from {file_path}")

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((48, 48)),
    transforms.ToTensor(),
])

def predict_single_image(model, image_path):
    # 加载并预处理图像
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)  # 增加一个维度以适应批处理格式
    
    # 将图像移到设备上
    image = image.to(device)
    
    # 模型预测
    with torch.no_grad():
        output = model(image)
        _, predicted = torch.max(output.data, 1)
    
    return predicted.item()

# 实例化模型
model = simpleconv3(num_classes=55)

# 加载模型权重
load_model(model, 'final_model.pth')

# 对单个图像进行预测
image_path = '"D:\Administrator\Desktop\all\all\1\img_332.jpg"'
prediction = predict_single_image(model, image_path)
print(f'Predicted class for the image: {prediction}')

运行该文件可对单图像进行预测