基于lab算法将黑白照片彩色化的评价模型代码

以下是一个基于 LAB 算法将黑白照片彩色化的评价模型的代码。该代码使用 PyTorch 框架实现。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from PIL import Image
import numpy as np

# 定义评价模型的网络结构
class ColorizationNet(nn.Module):
    def __init__(self):
        super(ColorizationNet, self).__init__()
        # 输入层
        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU(inplace=True)
        # 下采样
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.relu3 = nn.ReLU(inplace=True)
        self.conv4 = nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1)
        self.relu4 = nn.ReLU(inplace=True)
        self.conv5 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
        self.relu5 = nn.ReLU(inplace=True)
        self.conv6 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
        self.relu6 = nn.ReLU(inplace=True)
        self.conv7 = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1)
        self.relu7 = nn.ReLU(inplace=True)
        self.conv8 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
        self.relu8 = nn.ReLU(inplace=True)
        self.conv9 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
        self.relu9 = nn.ReLU(inplace=True)
        self.conv10 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1)
        self.relu10 = nn.ReLU(inplace=True)
        self.conv11 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
        self.relu11 = nn.ReLU(inplace=True)
        self.conv12 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
        self.relu12 = nn.ReLU(inplace=True)
        self.conv13 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
        self.relu13 = nn.ReLU(inplace=True)

        # 反卷积层
        self.conv14 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1)
        self.relu14 = nn.ReLU(inplace=True)
        self.conv15 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
        self.relu15 = nn.ReLU(inplace=True)
        self.conv16 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
        self.relu16 = nn.ReLU(inplace=True)

        # 上采样
        self.conv17 = nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1)
        self.relu17 = nn.ReLU(inplace=True)
        self.conv18 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
        self.relu18 = nn.ReLU(inplace=True)

        self.conv19 = nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1)
        self.relu19 = nn.ReLU(inplace=True)
        self.conv20 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
        self.relu20 = nn.ReLU(inplace=True)
        self.conv21 = nn.Conv2d(64, 2, kernel_size=3, stride=1, padding=1)

    def forward(self, gray):
        # 编码器
        x1 = self.conv1(gray)
        x2 = self.relu1(x1)
        x3 = self.conv2(x2)
        x4 = self.relu2(x3)
        x5 = self.conv3(x4)
        x6 = self.relu3(x5)
        x7 = self.conv4(x6)
        x8 = self.relu4(x7)
        x9 = self.conv5(x8)
        x10 = self.relu5(x9)
        x11 = self.conv6(x10)
        x12 = self.relu6(x11)
        x13 = self.conv7(x12)
        x14 = self.relu7(x13)
        x15 = self.conv8(x14)
        x16 = self.relu8(x15)
        x17 = self.conv9(x16)
        x18 = self.relu9(x17)
        x19 = self.conv10(x18)
        x20 = self.relu10(x19)
        x21 = self.conv11(x20)
        x22 = self.relu11(x21)
        x23 = self.conv12(x22)
        x24 = self.relu12(x23)
        x25 = self.conv13(x24)
        x26 = self.relu13(x25)

        # 解码器
        x27 = self.conv14(x26)
        x28 = self.relu14(x27)
        x29 = self.conv15(x28)
        x30 = self.relu15(x29)
        x31 = self.conv16(x30)
        x32 = self.relu16(x31)

        x33 = self.conv17(x32)
        x34 = self.relu17(x33)
        x35 = self.conv18(x34)
        x36 = self.relu18(x35)

        x37 = self.conv19(x36)
        x38 = self.relu19(x37)
        x39 = self.conv20(x38)
        x40 = self.relu20(x39)
        x41 = self.conv21(x40)

        return x41


# 定义评价模型
class ColorizationModel:
    def __init__(self, model_path):
        self.net = ColorizationNet()
        self.net.load_state_dict(torch.load(model_path))
        self.net.eval()

        # 定义预处理函数
        self.preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(256),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5], std=[0.5])
        ])

    def evaluate(self, gray_path, color_path):
        # 加载图片
        gray_img = Image.open(gray_path).convert('L')
        color_img = Image.open(color_path).convert('RGB')

        # 预处理
        gray_tensor = self.preprocess(gray_img)
        color_tensor = self.preprocess(color_img)

        # 扩展维度
        gray_tensor = gray_tensor.unsqueeze(0)
        color_tensor = color_tensor.unsqueeze(0)

        # 生成彩色图片
        with torch.no_grad():
            pred = self.net(gray_tensor)
            pred = torch.tanh(pred)
            pred = pred / 2 + 0.5
            pred = pred.squeeze(0).permute(1, 2, 0).numpy()

        # 计算 PSNR 和 SSIM
        pred = pred * 255
        pred = pred.astype(np.uint8)
        color_img = np.array(color_img)
        psnr = self.calculate_psnr(pred, color_img)
        ssim = self.calculate_ssim(pred, color_img)

        return psnr, ssim

    # 计算 PSNR
    def calculate_psnr(self, img1, img2):
        mse = np.mean((img1 - img2) ** 2)
        if mse == 0:
            return float('inf')
        max_pixel = 255.0
        psnr = 20 * np.log10(max_pixel / np.sqrt(mse))
        return psnr

    # 计算 SSIM
    def calculate_ssim(self, img1, img2):
        C1 = (0.01 * 255) ** 2
        C2 = (0.03 * 255) ** 2
        img1 = img1.astype(np.float64)
        img2 = img2.astype(np.float64)
        kernel = cv2.getGaussianKernel(11, 1.5)
        window = np.outer(kernel, kernel.transpose())

        mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
        mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
        mu1_sq = mu1 ** 2
        mu2_sq = mu2 ** 2
        mu1_mu2 = mu1 * mu2
        sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
        sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
        sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
        return ssim_map.mean()

代码中的 `ColorizationNet` 定义了彩色化网络的结构，包括编码器和解码器。`ColorizationModel` 定义了评价模型，包括加载网络模型和计算 PSNR 和 SSIM 的方法。其中 `calculate_psnr` 和 `calculate_ssim` 方法分别实现了计算 PSNR 和 SSIM 的算法。在 `evaluate` 方法中，我们首先加载灰度图和彩色图，并对其进行预处理。然后使用网络对灰度图进行彩色化，并计算生成的彩色图与原图的 PSNR 和 SSIM。最后返回这两个评价指标的值。