水下图像增强(论文复现)

本文所涉及所有资源均在 传知代码平台 可获取。
目录

概述

一、论文思路

二、模型介绍:

三、实现方法

四、复现过程(重要)

部署方式

概述

2021年11月,提出一种用于水下图像增强的U型Transformer模型,这是首次在水下图像增强任务中使用Transfomer模型,并且作者同时也发布了《U-shape Transformer for Underwater Image Enhancement》这篇文章。它主要针对水下图像增强任务,通过神经网络训练的方式,将模糊的,低分辨率的,对比度低的水下图像,转换成高清的、高分辨率的,对比度高的图像。并且作者也发布了一个大型的水下图像数据集LSUI,为后续在水下图像增强方向提供重要贡献。

一、论文思路

水下杂质的光吸收和散射导致水下成像质量差。现有的基于数据驱动的水下图像增强(UIE)技术缺乏包含各种水下场景和高保真参考图像的大规模数据集。此外,没有充分考虑不同颜色通道和空间区域的不一致衰减。为此,本文建立了大尺度水下图像(LSUI)数据集,并设计了U型Transformer。U型变压器集成了通道型多尺度特征融合变压器(CMSFFT)模块和针对UIE任务设计的空间型全局特征建模变压器(SGFMT)模块,增强了网络对衰减更严重的颜色通道和空间区域的关注。

二、模型介绍:

整体架构:包括基于 CMSFFT 和 SGFMT 的生成器和鉴别器。在生成器中,

(1) 编码:除了直接输入到网络之外,原始图像将分别进行3次下采样。然后经过11卷积后,将三个尺度特征图输入到对应的尺度卷积块中。四个卷积块的输出是CMSFFT和SGFMT的输入。

(2) 解码: 特征重新映射后,SGFMT输出直接发送到第一个卷积块。同时,4个卷积不同规模的区块将接收来自CMSFFT的四个输出。在判别器中,四个卷积块的输入包括:自身上层输出的特征图、来自解码部分的相应尺寸的特征图以及下采样到相应尺寸后通过11卷积生成的特征图使用参考图像。通过所描述的多尺度连接,梯度流可以在生成器和鉴别器之间在多个尺度上自由流动,从而可以获得稳定的训练过程,丰富生成图像的细节。

三、实现方法

1、SGFMT:用于替代生成器原来的瓶颈层,可以辅助网络对全局信息进行建模,并加强网络对严重退化部分的关注。具体流程大概为:输入特征图通过线性投影转化为一维序列然后嵌入位置编码进入Transformer层(每一个Transformer层包含一个多头注意力块(MHA)和一个前馈神经网络(FFN)),得到输出序列后通过特征映射转化为输出特征图。

2、CMSFFT: 为了加强网络对衰减更严重的颜色通道的关注,我们设计了CMSFFT块来代替原始生成器编码解码架构的跳跃连接,它由以下三个部分组成。

(1)、多尺度特征编码(Multi-Scale Feature Encoding):输入是不同尺度的特征图, 直接应用于分区原始图像的线性投影不同,我们使用具有相关滤波器大小在不同尺度的特征图上进行线性投影。

(2)、通道方式多头注意力(CMHA):CMHA 块有 6 个输入,其中 IN 表示实例规范化操作。这种注意力操作沿着通道轴而不是经典的补丁轴进行,可以引导网络关注图像质量下降更严重的通道。此外,在相似度图上使用IN来帮助梯度流顺利传播。

(3)、前馈网络(FFN):与前向传播类似,包含多层感知机(MLP)和归一化层(LN)。

3、损失函数:为了利用LAB和LCH颜色空间更宽的色域表示范围以及更准确地描述颜色饱和度和亮度,我们设计了结合RGB、LAB和LCH颜色空间的多颜色空间损失函数来训练我们的网络。 其中α、β、γ、μ为超参数,经过大量实验分别设置为0.001、1、0.1、100。后面的四个分别是生成器的损失函数它们是定义好的。

四、复现过程(重要)

先看结果,原图像与增强之后的图像对比

能非常直观的感觉出增强效果还是非常好的,接下来是具体的步骤。

代码结构

1、在Pycharm中导入项目;

2、下载数据集LSUI并将数据集添加到项目data目录中

数据集下载链接:详见附件;

3、修改保存权重路径和数据集路径

(1)、修改数据集路径

在mytrain.py

在mytest.py

(2)、修改权重路径

mytest.py

4、测试之前预训练的模型

预训练模型链接:详见附件

可以直接使用笔者之前训练好的权重去直接测试mytest.py

5、运行mytrain.py文件:

6、运行test.ipynb文件:

7、得出输出图像

部署方式

python 3.7, Pytorch 1.8

编译器采用Pycharm,拿到代码之后,结合ReadMe以及"requirements.txt"配置好环境之后,可以直接使用预训练的模型去处理水下图像;也可以根据自己的需求重新训练一整个网络模型。

感觉不错,点击我,立即使用

相关推荐
lzl204011 分钟前
【深度学习总结】热力图-Grad-CAM使用
人工智能·深度学习·热力图
肖遥Janic26 分钟前
Stable Diffusion绘画 | 签名、字体、Logo设计
人工智能·ai·ai作画·stable diffusion
新手小白勇闯新世界31 分钟前
机器学习笔记(三)-模型评估与选择后半部分
人工智能·笔记·机器学习
Nil_cxc37 分钟前
机器学习周报(9.30-10.6)
人工智能·机器学习
我要吐泡泡了哦41 分钟前
GAMES104:16 游戏引擎的玩法系统:基础AI-学习笔记
人工智能·学习·游戏引擎
别NULL1 小时前
毕业设计项目——基于频域分析的MP3格式音频文件到MIDI格式转换(论文/代码)
人工智能·算法·课程设计
是Winky啊1 小时前
【论文阅读】超分辨率图像重建算法综述
人工智能·重构
林学长@F5 AI社区2 小时前
【AI人脸融合】FaceFusion 3.0一键整合包,支持AI视频/图片人脸融合、数字人视频、表情修复、年龄修改等
人工智能·音视频
美狐美颜sdk2 小时前
如何搭建直播美颜平台?视频美颜SDK的核心技术详解
人工智能·音视频·直播美颜sdk·第三方美颜sdk·美狐美颜sdk
白葵新2 小时前
PCL 3D-SIFT关键点检测(曲率不变特征约束
c++·人工智能·算法·计算机视觉·3d