提示:海水梦悠悠,君愁我亦愁。

文章目录


这里若是你要写论文用到这篇文章,你使用的技术主要包括深度学习框架PyTorch ,Web框架:Flask,图像处理:PIL,Numpy,机器学习算法:KNN,前端:HTML,CSS,JavaScript。

通过深度学习Step1这个文章来看,模型其实去噪能力还不是那么强,所以若是我们要提高去噪能力也就需要提高模型的复杂度,而提高模型复杂度的方式通常是只有两种,一种是提高模型的特征的维度,还一种是提高模型的深度,这两个到底哪个更好,这个不一定,,通常是深度越深则效率越高,因为深度越深则可以通过更少的参数起到一样的效果,还有就是上一篇只是使用了一个图片进行模型的预测,所以模型的泛化能力肯定是很弱,
所以这一篇写的就是就是读取多个图片提高模型泛化能力的一个项目,并且之前是在juypter上一步一步执行的,并且上一个是在同一个文件中,所以这一篇也将上一个的分开,放入到多个文件夹中

第四章 构建文件

将要处理的大量数据放在common目录下,数据包是dataset
在image_denoising包下创建以下文件夹
config.py 项目配置
data.py 数据处理
model.py 模型定义模块
train.py 数据加载,预处理,

按照上一个Step1中流程,需要先读取数据,这里将全部数据进行读出,划分数据集我们放在train中
在这里插入图片描述
首先先定义config.py文件

import os

import torch

from common.utils import seed_everything

current_dir = os.getcwd()
parent_dir = os.path.dirname(current_dir)# 上一级目录
IMG_HEIGHT = 256
IMG_WIDTH = 256 # 图片尺寸
SEED = 42 # 随机种子
seed_everything(SEED)
BATCH_SIZE = 128
Train_ratio = 0.7 # 训练集比例
VAL_ratio = 0.8
Noise_level = 0.1 # 噪声等级
Epochs = 1
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Learning_rate = 1e-3 # 学习率

定义data.py 文件

import os

from PIL import Image
from torchvision import transforms
from torch.utils.data import Dataset
import re
from image_denoising.config import parent_dir
# 定义数据集类型
class MyDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        # img_dir:指定图像文件所在的目录路径,用于定位和读取图像数据。
        # transform:图像变换操作,用于对图像进行预处理(如缩放、裁剪、转换为张量等)。当值为 None 时,表示不进行任何变换处理,读取图片因为要进行维度的转化所以需要这个参数
        self.img_dir = img_dir
        self.transform = transform
        self.img_name_list = sorted(os.listdir(img_dir), key=lambda x: int(re.search(r'\d+', x).group()))
        print(self.img_name_list)



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

    def __getitem__(self, idx):
        # 先找文件的完整路径
        img_name = self.img_name_list[idx]
        img_path = os.path.join(self.img_dir, img_name)
        # 读取图片,但是要注意图片的格式,使用Image读取图片的时候默认是RGBA
        # 这里我使用的是RGB格式的图片,所以这里需要转化为RGB格式
        img = Image.open(img_path).convert('RGB')
        if self.transform:
            img = self.transform(img)
        else:
            raise ValueError('transform is None')
        return img  # 返回原始图片


if __name__ == '__main__':
    # 测试数据集排序是否正确
    img_dir = '../common/dataset'
    print(os.listdir(img_dir))
    img_name_list = sorted(os.listdir(img_dir), key=lambda x: int(re.search(r'\d+', x).group()))
    print(img_name_list)

    #测试此模块的代码是否正确
    dataset = MyDataset(img_dir)
    # 接下来依然是和上一篇文章一样,定义transform
    transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
    ])
    # 创建数据集对象
    dataset = MyDataset(img_dir=parent_dir + '/common/dataset/',
                        transform=transform)  # 注意这里需要多一个斜杠,因为是直接将这个地址和文件的名字直接进行拼接得到特征
    print(len(dataset))

依然按照Step1中数据好了之后,开始构建模型model.py

# 创建模型,这里依然是直接使用上一个步骤的模型
import torch
import torch.nn as nn


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.encoder = nn.Sequential(
            # 第一层卷积池化层
            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),  # 第一层卷积层
            nn.ReLU(),  # 激活函数(激活函数的作用简单来说从A'到A的过程肯定是一个函数存在的,而激活函数的作用就是让这个函数更加多样化)
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
            nn.ReLU(),
            # 第二层卷积池化层
            nn.Conv2d(16, 8, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
            nn.ReLU()
        )
        self.decoder = nn.Sequential(
            # 第一层转置卷积层
            nn.ConvTranspose2d(8, 16, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            # 第二层转置卷积层
            nn.ConvTranspose2d(16, 3, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid(),  # 这里因为输入是0~1的数,所以输出也是0~1的数,这里因为是最后一层,所以使用Sigmoid激活函数
        )

    def forward(self, x):  # 前向传播,也就相当于是完整的走一遍这个模型的流程
        x = self.encoder(x)
        x = self.decoder(x)
        return x

if __name__ == '__main__':
    # 测试模型是否正确
    input=torch.randn(1, 3, 256, 256)
    model = MyModel()
    output = model(input)
    print(output.shape)

模型定义完毕之后接下来就是训练测试,

from torch import nn
# 创建训练集,测试集
from torch.utils.data import random_split, Subset  # 这个包的作用是划分数据集
from torchvision import transforms
from tqdm import tqdm

from image_denoising.data import MyDataset
from image_denoising.model import MyModel
from image_denoising.config import device, Epochs, BATCH_SIZE
import matplotlib
matplotlib.use('Agg')  # 添加这一行在导入 pyplot 之前
import matplotlib.pyplot as plt  # 引入这个模块的作用是为了看到这个图片
# 创建数据加载器
import torch
from torch.utils.data import DataLoader
from image_denoising.config import parent_dir

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])
dataset = MyDataset(img_dir=parent_dir + '/common/dataset/',
                    transform=transform)  # 注意这里需要多一个斜杠,因为是直接将这个地址和文件的名字直接进行拼接得到特征
# train_set, val_set = random_split(dataset, [0.7, 0.3])  # 将数据集进行划分

# 然后替换原来的 dataset 使用 partial_dataset
train_set, val_set = random_split(dataset, [0.7, 0.3])

train_loader = torch.utils.data.DataLoader(  # 分批次加载训练数据
    train_set,
    batch_size=BATCH_SIZE,
    shuffle=True,
    drop_last=True  # 如果数据集大小不能被batch_size整除,则最后一个batch将被丢弃
)
test_loader = torch.utils.data.DataLoader(  # 分批次加载测试数据
    val_set,
    batch_size=BATCH_SIZE,
    shuffle=False,
)

model = MyModel()
# 定义损失函数和优化器
loss = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 训练模型
min_loss = float('inf')
model.to(device)
for epoch in tqdm(range(1, Epochs + 1)):
    model.train()
    train_loss = 999
    for data in train_loader:
        data = data.to(device)
        inputs = data + torch.randn_like(data)  # data是真实数据,我们使用噪音数据进行训练,真实数据用于计算loss
        imputs=torch.clamp(inputs, min=0, max=1)# 使得数据处于0~1这个区间,
        outputs = model(inputs)
        loss_value = loss(outputs, data)
        loss_value.backward()
        optimizer.step()
        optimizer.zero_grad()
        train_loss += loss_value.item()
    if train_loss < min_loss:
        min_loss = train_loss
        torch.save(model.state_dict(), 'model.pth')
        print(f'Save model with loss: {min_loss:.4f}')
# 测试并画图
model.eval()  # 将模型设置为评估模式,关闭dropout和batch normalization的训练行为

test_loss = 0  # 初始化测试损失值
with torch.no_grad():  # 在评估阶段不需要计算梯度,节省内存并加快计算
    for data in test_loader:  # 遍历验证数据加载器中的每个批次
        data = data.to(device)  # 将数据移动到指定设备(CPU或GPU)
        #加载最优模型
        model.load_state_dict(torch.load('model.pth'))
        outputs = model(data)  # 前向传播,获取模型输出

        # 计算并打印输出张量的形状,便于调试和确认维度
        print(outputs.shape)

			
        # 将模型输出转换为numpy数组以便可视化,下面的操作是将原始图,加噪音图,去噪之后的图呈现处理,因为涉及到三个图并排输出的问题
        #所以需要使用zip,对于zip不了解的同学,,不过下面的我已经解释的比较详细了,看看各位能不能看懂,可以搜索一下,
        # detach(): 断开与计算图的连接
        # permute(0, 2, 3, 1): 将通道维度从第1位移到最后(CHW -> HWC)
        # cpu(): 移动到CPU上
        # numpy(): 转换为numpy数组
        img_numpy = outputs.detach().permute(0, 2, 3, 1).cpu().numpy()

        # 将原始干净图像转换为numpy数组用于对比显示
        img_data = data.detach().permute(0, 2, 3, 1).cpu().numpy()

        # 生成带噪声的图像用于对比显示
        # data + torch.randn_like(data): 给原始数据添加随机噪声,并且这个数据要在0~1之间,因为画图需要这个区间
        img_no = torch.clamp((data + torch.randn_like(data)), min=0, max=1).permute(0, 2, 3, 1).cpu().numpy()

        # 创建一个子图网格用于显示三组图像:噪声图像、原始图像和去噪后的图像
        # 3行1列的子图布局,图像尺寸为(10, 5)
        fig, axs = plt.subplots(3, 1, figsize=(10, 5))

        # 遍历三种类型的图像数据,并将它们分别显示在对应的行中
        # img_no: 带噪声的图像
        # img_data: 原始干净图像
        # img_numpy: 模型去噪后的图像
        for imgs, row in zip([img_no, img_data, img_numpy], axs):
            # 遍历每行中的各个子图位置
            for img, ax in zip(imgs, row):
                # 显示图像
                ax.imshow(img)
                # 关闭坐标轴显示,使图像更清晰
                ax.axis('off')

        plt.savefig(f'denoising_result.png')  # 保存图像
        plt.close()  # 关闭图像

最后运行呈现的是,我这里只训练了一轮,所以效果不是那么完美
在这里插入图片描述

简单来说就是和之前一样,不过这里就是读取数据大了,所以需要批次读取,并且这里因为数据处理较大,所以将模型和数据移动到了GPU上进行处理,其他来看和上一篇文章深度学习Step1数据处理的过程差不多,拿到数据,定义模型,计算损失,进行测试,画图,相信你将上一个理解了再来理解这个就很简单,这个相当于是上一篇的进阶,最后今天的全部代码包请关注公众号“Sun小明同学获取”

更多推荐