基于GAN的水下图像增强系统设计与实现

Confusd_pai

已于 2024-12-13 11:35:14 修改

阅读量2.1k

点赞数 35

分类专栏：项目文章标签：生成对抗网络人工智能神经网络

于 2024-12-13 11:07:42 首次发布

本文链接：https://blog.csdn.net/Confusd_pai/article/details/144444875

版权

项目专栏收录该内容

6 篇文章

订阅专栏

（源码【已上传到github】：https://github.com/0101110AC/My-project/tree/main）

1. 系统概述

本系统是一个基于生成对抗网络（GAN）的水下图像增强解决方案，旨在提升水下图像的质量，包括去噪、颜色增强和细节恢复。系统由训练和测试两大模块组成，采用模块化设计，具有良好的扩展性和维护性。

1.1模型概述

本文提出了一个针对不同复杂场景量身定制的水下图像增强模型。为了实现这个目标，本文在现有的物理模型的基础上，加入了机器学习和深度学习的框架，以便于提升模型的鲁棒性和适应性。

首先，该模型使用了U-Net类似的编码器-解码器结构，U-Net网络结构是一个对称结构，如图1所示。蓝/白色框表示 feature map；蓝色箭头表示 3x3 卷积，用于特征提取；灰色箭头表示 skip-connection，用于特征融合；红色箭头表示池化 pooling，用于降低维度；绿色箭头表示上采样 upsample，用于恢复维度；青色箭头表示 1x1 卷积，用于输出结果。

图1. U-Net结构详解图

模型在使用了U-Net类似的编码器-解码器结构的基础上，采用GAN架构进行对抗训练，具体如图2所示。相比于传统的方法， GAN架构可以生成高质量的新数据，处理复杂的数据分布。在特征学习上，GAN可以自动提取特征，学习到更具代表性的特征。GAN适用于多种任务和领域，具有非常强的灵活性。因此，本文选用GAN架构代替传统的方法完成工作。

图2. GAN架构

之所以选择U-Net和GAN的组合，是因为这种架构特别适合图像增强任务。这个模型能保留原始图像的细节信息，能够学习复杂的图像转换，可以通过对抗训练提高生成质量，对于像素级损失确保基本内容保持不变。

1.2评估指标概述

本文通过计算PSNR（峰值信噪比）、UCIQE（颜色质量指数）和UIQM（视觉质量指数）来计算图像质量的评估指标。

计算PSNR。本文计算了原始图像和增强图像之间的均方误差（MSE），然后根据均方误差计算了PSNR。PSNR是一种常用的图像质量评估指标，它反映了图像的清晰度和细节保留程度。计算的值越高，说明图像质量越好。

计算UCIQE。本文将图像从RGB颜色空间转换为LAB颜色空间，然后提取亮度通道（L通道）和色度通道（A和B通道）。接着，计算了色度通道的均值、标准差和最大值，以及亮度通道的最大值和最小值。最后根据这些值计算了UCIQE。UCIQE是一种常用的图像质量评估指标，它反映了图像的颜色质量和对比度。计算的值越高，说明图像质量越好。

计算UIQM。本文将图像从RGB颜色空间转换为HSV颜色空间，然后提取色度通道（S通道）和亮度通道（V通道）。接着计算色度通道的均值、拉普拉斯算子的方差和标准差。最后根据这些值计算了UIQM。UIQM是一种常用的图像质量评估指标，它反映了图像的视觉质量。计算的值越高，说明图像质量越好。

这些指标互相补充，共同提供了对水下图像质量的全面评估。

2. 核心架构设计

2.1 生成器（Generator）

生成器采用编码器-解码器架构，结合多种先进的深度学习技术：

2.1.1. 特征提取与编解码

class Generator(nn.Module):

    def __init__(self):

        # 初始特征提取

        self.init_conv = ConvBlock(3, 64)

        

        # 编码器

        self.encoder = nn.Sequential(

            ConvBlock(64, 128, 4, 2, 1),

            ConvBlock(128, 256, 4, 2, 1),

        )

        

        # 解码器

        self.decoder = nn.Sequential(

            nn.ConvTranspose2d(256, 128, 4, 2, 1),

            nn.BatchNorm2d(128),

            nn.ReLU(inplace=True),

            # ...

        )

2.1.2. 注意力机制

# 空间注意力

class SpatialAttention(nn.Module):

    def forward(self, x):

        avg_out = torch.mean(x, dim=1, keepdim=True)

        max_out, _ = torch.max(x, dim=1, keepdim=True)

        x_cat = torch.cat([avg_out, max_out], dim=1)

        attention = self.sigmoid(self.conv(x_cat))

        return x * attention



# 通道注意力

class ChannelAttention(nn.Module):

    def forward(self, x):

        avg_out = self.fc(self.avg_pool(x))

        max_out = self.fc(self.max_pool(x))

        attention = self.sigmoid(avg_out + max_out)

        return x * attention

2.1.3. 图像增强处理

def enhance_colors(self, img):

    """HSV空间的颜色增强"""

    img_hsv = torch.zeros_like(img)

    img_hsv[:, 0] = 0.5 * (torch.sin(3.14159 * img[:, 0]) + 1)

    img_hsv[:, 1] = torch.clamp(1.2 * img[:, 1], 0, 1)

    img_hsv[:, 2] = torch.clamp(1.1 * img[:, 2], 0, 1)

    return img_hsv



def denoise(self, img, kernel_size=3, sigma=1.0):

    """高斯滤波去噪"""

    # 实现高斯核和滤波操作

2.2 判别器（Discriminator）

判别器采用经典的CNN结构，通过多层卷积实现图像真伪判断：

class Discriminator(nn.Module):

    def __init__(self):

        self.model = nn.Sequential(

            nn.Conv2d(3, 64, 4, stride=2, padding=1),

            nn.LeakyReLU(0.2, inplace=True),

            # 逐步降采样

            # ...

            nn.Sigmoid()

        )

3. 数据处理与训练

3.1 数据集实现

class UnderwaterDataset(Dataset):

    def __init__(self, root_dir, transform=None):

        self.images_dir = Path(root_dir) / 'images'

        self.transform = transform

        self.images = list(self.images_dir.glob('*.jpg'))

3.2 训练配置

class Config:

    def __init__(self):

        self.epochs = 200

        self.batch_size = 8

        self.lr = 0.0002

        self.b1 = 0.5

        self.b2 = 0.999

        self.image_size = 256

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

4. 测试与评估系统

4.1 测试器设计

class Tester:

    def __init__(self):

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self.image_size = 256

        self.model_path = self.find_latest_model()

4.2 评估指标实现

4.2.1. PSNR（峰值信噪比）

def calculate_psnr(self, original, enhanced):

    return psnr(original, enhanced)

4.2.2. UCIQE（水下图像彩色质量评估）

def calculate_uciqe(self, img):

    lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)

    l, a, b = cv2.split(lab)

    chroma = np.sqrt(np.square(a) + np.square(b))

    

    uc = np.mean(chroma)

    sc = np.sqrt(np.mean(np.square(chroma - uc)))

    cont = np.max(l) - np.min(l)

    

    return 0.4680 * sc + 0.2745 * cont + 0.2576 * uc

4.2.3. UIQM（水下图像质量度量）

def calculate_uiqm(self, img):

    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)

    h, s, v = cv2.split(hsv)

    

    uicm = np.mean(s)

    uism = cv2.Laplacian(v, cv2.CV_64F).var()

    uiconm = np.std(v)

    

    return 0.0282 * uicm + 0.2953 * uism + 0.6765 * uiconm