pytorch: 影像恢復問題的代碼實作詳解(derain,dehaze,deblur,denoise等通用)-有解無憂

文章目錄

- - 前言
  - 資料集
  - - 訓練資料集
    - 評估資料集
    - 測驗資料集
  - 網路模型
  - 自定義工具包
  - 網路訓練和測驗
  - 結語

前言

影像恢復是一類圖形去噪問題的集合，在深度學習中可以理解為監督回歸問題，主要包括影像去雨、影像去霧、影像去噪，影像去模糊和影像去馬賽克等內容，但利用 pytorch 實作的代碼類似，只是在具體網路結構上略有區別，

以影像去雨為例，之前寫過一篇影像去雨的 pytorch 實作文章: https://blog.csdn.net/Wenyuanbo/article/details/116541682，但因當時能力和水平有限，實作邏輯存在問題，最近重新整理分享一下，希望能對大家有所幫助，工程檔案如圖所示，資料集路徑根據自己情況設定，

在這里插入圖片描述
、

資料集

利用監督回歸方法實作影像去雨時，一般資料集為有雨圖和無雨圖成對存在，首先我喜歡習慣性的將所有成對資料分別從 0 到結束對應重新排序(這個其實不影響，具體自己設計即可)，

MyDataset.py

import os
import random
import torchvision.transforms.functional as ttf
from torch.utils.data import Dataset
from PIL import Image

訓練資料集

訓練資料集是用來整合訓練資料的，將有雨圖和無雨圖分別對應進行剪切，轉張量等操作，

class MyTrainDataSet(Dataset):  # 訓練資料集
    def __init__(self, inputPathTrain, targetPathTrain, patch_size=128):
        super(MyTrainDataSet, self).__init__()

        self.inputPath = inputPathTrain
        self.inputImages = os.listdir(inputPathTrain)  # 輸入圖片路徑下的所有檔案名串列

        self.targetPath = targetPathTrain
        self.targetImages = os.listdir(targetPathTrain)  # 目標圖片路徑下的所有檔案名串列

        self.ps = patch_size

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

    def __getitem__(self, index):

        ps = self.ps
        index = index % len(self.targetImages)

        inputImagePath = os.path.join(self.inputPath, self.inputImages[index])  # 圖片完整路徑
        inputImage = Image.open(inputImagePath).convert('RGB')  # 讀取圖片

        targetImagePath = os.path.join(self.targetPath, self.targetImages[index])
        targetImage = Image.open(targetImagePath).convert('RGB')

        inputImage = ttf.to_tensor(inputImage)  # 將圖片轉為張量
        targetImage = ttf.to_tensor(targetImage)

        hh, ww = targetImage.shape[1], targetImage.shape[2]  # 圖片的高和寬

        rr = random.randint(0, hh-ps)  # 亂數： patch 左下角的坐標 (rr, cc)
        cc = random.randint(0, ww-ps)
        # aug = random.randint(0, 8)  # 亂數，對應對圖片進行的操作

        input_ = inputImage[:, rr:rr+ps, cc:cc+ps]  # 裁剪 patch ，輸入和目標 patch 要對應相同
        target = targetImage[:, rr:rr+ps, cc:cc+ps]

        return input_, target

評估資料集

在網路訓練中，不一定最后一次訓練的效果就是最好的，評估資料集是在每一個 epoch 訓練結束后對網路訓練的性能進行評估，目的在于將最好的一次訓練結果保存，

class MyValueDataSet(Dataset):  # 評估資料集
    def __init__(self, inputPathTrain, targetPathTrain, patch_size=128):
        super(MyValueDataSet, self).__init__()

        self.inputPath = inputPathTrain
        self.inputImages = os.listdir(inputPathTrain)  # 輸入圖片路徑下的所有檔案名串列

        self.targetPath = targetPathTrain
        self.targetImages = os.listdir(targetPathTrain)  # 目標圖片路徑下的所有檔案名串列

        self.ps = patch_size

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

    def __getitem__(self, index):

        ps = self.ps
        index = index % len(self.targetImages)

        inputImagePath = os.path.join(self.inputPath, self.inputImages[index])  # 圖片完整路徑
        inputImage = Image.open(inputImagePath).convert('RGB')  # 讀取圖片,灰度圖

        targetImagePath = os.path.join(self.targetPath, self.targetImages[index])
        targetImage = Image.open(targetImagePath).convert('RGB')

        inputImage = ttf.center_crop(inputImage, (ps, ps))
        targetImage = ttf.center_crop(targetImage, (ps, ps))

        input_ = ttf.to_tensor(inputImage)  # 將圖片轉為張量
        target = ttf.to_tensor(targetImage)

        return input_, target

測驗資料集

測驗資料集的目的是將輸入有雨進行去雨得到去雨后的結果，注意輸入一般是原圖大小，不進行裁剪，

class MyTestDataSet(Dataset):  # 測驗資料集
    def __init__(self, inputPathTest):
        super(MyTestDataSet, self).__init__()

        self.inputPath = inputPathTest
        self.inputImages = os.listdir(inputPathTest)  # 輸入圖片路徑下的所有檔案名串列

    def __len__(self):
        return len(self.inputImages)  # 路徑里的圖片數量

    def __getitem__(self, index):
        index = index % len(self.inputImages)

        inputImagePath = os.path.join(self.inputPath, self.inputImages[index])  # 圖片完整路徑
        inputImage = Image.open(inputImagePath).convert('RGB')  # 讀取圖片

        input_ = ttf.to_tensor(inputImage)  # 將圖片轉為張量

        return input_

網路模型

以一個 5 層簡單卷積神經網路為例子，具體網路自己設定，
NetModel.py

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.inconv = nn.Sequential(  # 輸入層網路
            nn.Conv2d(3, 32, 3, 1, 1),
            nn.ReLU(inplace=True)
        )
        self.midconv = nn.Sequential(  # 中間層網路
            nn.Conv2d(3, 32, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(3, 32, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(3, 32, 3, 1, 1),
            nn.ReLU(inplace=True),
        )
        self.outconv = nn.Sequential(  # 輸出層網路
            nn.Conv2d(3, 32, 3, 1, 1),
        )
        
    def forward(self, x):

        x = self.inconv(x)
        x = self.midconv(x)
        x = self.outconv(x)
        
        return x

自定義工具包

自定義工具包主要是一個計算峰值信噪比(PSNR)的方法用來對訓練進行評估，

utils.py

import torch

def torchPSNR(tar_img, prd_img):
    imdff = torch.clamp(prd_img,0,1) - torch.clamp(tar_img,0,1)
    rmse = (imdff**2).mean().sqrt()
    ps = 20*torch.log10(1/rmse)
    return ps

網路訓練和測驗

main.py

import sys
import time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from tqdm import tqdm, trange  # 進度條
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from torch.autograd import Variable
from torch.optim.lr_scheduler import MultiStepLR
import utils
from NetModel import Net
from MyDataset import *

if __name__ == '__main__':  # 只有在 main 中才能開多執行緒
	EPOCH = 100  # 訓練次數
    BATCH_SIZE = 18  # 每批的訓練數量
    LEARNING_RATE = 1e-3  # 學習率
	
	inputPathTrain = 'Rain100H/inputTrain/'  # 訓練輸入圖片路徑
    targetPathTrain = 'Rain100H/targetTrain/'  # 訓練目標圖片路徑
    inputPathTest = 'Rain100H/inputTest/'  # 測驗輸入圖片路徑
    resultPathTest = 'Rain100H/resultTest/'  # 測驗結果圖片路徑
    middleResultPathTest = 'Rain100H/resultMiddleTest/'  # 測驗結果圖片路徑
    targetPathTest = 'Rain100H/targetTest/'  # 測驗目標圖片路徑

	myNet = Net()  # 實體化網路
    myNet = myNet.cuda()  # 網路放入GPU中
	
	criterion = nn.MSELoss().cuda()
	
    optimizer = optim.Adam(myNet.parameters(), lr=LEARNING_RATE)  # 網路引數優化演算法

    scheduler = MultiStepLR(optimizer, milestones=[30, 50, 80], gamma=0.2)  # 階段調整學習率

    # 訓練資料
    datasetTrain = MyTrainDataSet(inputPathTrain, targetPathTrain)  # 實體化訓練資料集類
    # 可迭代資料加載器加載訓練資料
    trainLoader = DataLoader(dataset=datasetTrain, batch_size=BATCH_SIZE, shuffle=True, drop_last=False, num_workers=6, pin_memory=True)

    # 評估資料
    datasetValue = MyValueDataSet(inputPathTest, targetPathTest)  # 實體化評估資料集類
    valueLoader = DataLoader(dataset=datasetValue, batch_size=16, shuffle=True, drop_last=False, num_workers=6, pin_memory=True)

	# 測驗資料
    datasetTest = MyTestDataSet(inputPathTest)  # 實體化測驗資料集類
    # 可迭代資料加載器加載測驗資料
    testLoader = DataLoader(dataset=datasetTest, batch_size=1, shuffle=False, drop_last=False, num_workers=6, pin_memory=True)

	# 開始訓練
    print('-------------------------------------------------------------------------------------------------------')
    if os.path.exists('./model_best.pth'):  # 判斷是否預訓練
        myNet.load_state_dict(torch.load('./model_best.pth'))

    for epoch in range(EPOCH):
        myNet.train()  # 指定網路模型訓練狀態
        iters = tqdm(trainLoader, file=sys.stdout)  # 實體化 tqdm，自定義
        epochLoss = 0  # 每次訓練的損失
        timeStart = time.time()  # 每次訓練開始時間
        for index, (x, y) in enumerate(iters, 0):

            myNet.zero_grad()  # 模型引數梯度置0
            optimizer.zero_grad()  # 同上等效

            input_train, target = Variable(x).cuda(), Variable(y).cuda()  # 轉為可求導變數并放入 GPU
            output_train = myNet(input_train)  # 輸入網路，得到相應輸出

            loss = criterion(output_train, target) # 計算網路輸出與目標輸出的損失

            loss.backward()  # 反向傳播
            optimizer.step()  # 更新網路引數
            epochLoss += loss.item()  # 累計一次訓練的損失

            # 自定義進度條前綴
            iters.set_description('Training !!!  Epoch %d / %d,  Batch Loss %.6f' % (epoch+1, EPOCH, loss.item()))

        # 評估
        myNet.eval()
        psnr_val_rgb = []
        for index, (x, y) in enumerate(valueLoader, 0):
            input_, target_value = x.cuda(), y.cuda()
            with torch.no_grad():
                output_value = myNet(input_)
            for output_value, target_value in zip(output_value[0], target_value):
                psnr_val_rgb.append(psnr(output_value, target_value))

        psnr_val_rgb = torch.stack(psnr_val_rgb).mean().item()

        if psnr_val_rgb > best_psnr:
            best_psnr = psnr_val_rgb
            best_epoch = epoch
            torch.save(myNet.state_dict(), 'model_best.pth')

        loss_list.append(epochLoss)  # 插入每次訓練的損失值
        lr_list.append(scheduler.get_lr())
        scheduler.step(epoch)  # 更新學習率
        torch.save(myNet.state_dict(), 'model.pth')  # 每次訓練結束保存模型引數
        timeEnd = time.time()  # 每次訓練結束時間
        print("------------------------------------------------------------")
        print("Epoch:  {}  Finished,  Time:  {:.4f} s,  Loss:  {:.6f}.".format(epoch+1, timeEnd-timeStart, epochLoss))
        print('-------------------------------------------------------------------------------------------------------')
    print("Training Process Finished ! Best Epoch : {} , Best PSNR : {:.2f}".format(best_epoch, best_psnr))

	# 測驗
    print('--------------------------------------------------------------')
    myNet.load_state_dict(torch.load('./model_best.pth'))  # 加載已經訓練好的模型引數
    myNet.eval()  # 指定網路模型測驗狀態

    with torch.no_grad():  # 測驗階段不需要梯度
        timeStart = time.time()  # 測驗開始時間
        for index, x in enumerate(tqdm(testLoader, desc='Testing !!! ', file=sys.stdout), 0):
            torch.cuda.empty_cache()  # 釋放顯存
            input_test = x.cuda()  # 放入GPU
            output_test = myNet(input_test)  # 輸入網路，得到輸出
            save_image(output_test[0], middleResultPathTest + str(index+1).zfill(3) + tail)  # 保存網路輸出結果
            save_image(output_test[1], resultPathTest + str(index+1).zfill(3) + tail)  # 保存網路輸出結果
        timeEnd = time.time()  # 測驗結束時間
        print('---------------------------------------------------------')
        print("Testing Process Finished !!! Time: {:.4f} s".format(timeEnd - timeStart))

	# 繪制訓練時損失曲線
    plt.figure(1)
    x = range(0, EPOCH)
    plt.xlabel('epoch')
    plt.ylabel('epoch loss')
    plt.plot(x, loss_list, 'r-')
    # 繪制學習率改變曲線
    plt.figure(2)
    plt.xlabel('epoch')
    plt.ylabel('learning rate')
    plt.plot(x, lr_list, 'r-')

    plt.show()

結語

關于影像恢復特別是影像去雨問題歡迎一起交流學習，

轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/298046.html

標籤：AI

上一篇：文本分類煉丹實錄（上篇）

下一篇：【大局觀】前言——這個世界為什么越來越不平靜