当前位置: 首页 > AI > 文章内容页

【图像去噪】第六期论文复现赛——MIRNet

时间:2025-07-23    作者:游乐小编    

本文复现MIRNet系列论文,含V1和V2版本。V1先提取低级特征,经递归残差组处理,再得残差图像,最终恢复图像;V2类似但优化模块提升速度。复现精度达标,提供数据集、预训练模型、文件结构及训练、评估等操作方法,方便使用。

【图像去噪】第六期论文复现赛——mirnet - 游乐网

MIRNet 系列论文复现 兼 飞桨特色模型挑战赛

复现了两篇论文:

MIRNetV1: Learning Enriched Features for Real Image Restoration and Enhancement

MIRNetV2: Learning Enriched Features for Fast Image Restoration and Enhancement

最新源码:https://github.com/swz30/MIRNet 和 https://github.com/swz30/MIRNetV2

复现地址:https://github.com/sldyns/MIRNetV2_paddle

1. 简介

MIRNet V1:

【图像去噪】第六期论文复现赛——MIRNet - 游乐网        

给定一个图像 IRH×H×3I∈RH×H×3,MIRNet 首先应用一个卷积层来提取低级特征 X0RH×W×CX0∈RH×W×C. 接下来,特征映射 X0X0 通过 NN 个递归残差组(RRGs),产生深度特征 XdRH×W×CXd∈RH×W×C. 我们注意到每个 RRG 包含多个多尺度残差块(MRB),MRB 由多个(本文中有三个)并行连接的全卷积流组成,每个连接上先由 DAU 抑制了不太有用的特性,并且只允许信息更丰富的特性进一步传递给 SKFF,SKFF 模块通过 Fuse 和 select 两种操作对接受域进行动态调整. 接下来, 应用卷积层,得到残差图像 RRH×H×3R∈RH×H×3。最后,恢复的图像为 I^=I+RI^=I+R.

MIRNet V2:

【图像去噪】第六期论文复现赛——MIRNet - 游乐网        

MIRNetV2 过程同 MIRNet V1 类似,主要在 MRB 内减少了卷积流间的链接,同时将DAU替换为残差上下文块,显著降低了模型大小并提升了推理速度.

2. 复现精度

MIRNet V1:

验收标准:SIDD PSNR: 39.678

复现结果:SIDD PSNR: 39.687

飞桨特色模型挑战赛:

验收标准:SIDD PSNR: 37,SSIM: 0.94

MIRNet V2,训练92个epoch精度:SIDD PSNR: 39.5286,SSIM: 0.9578

3. 数据集、预训练模型、文件结构

3.1 数据集

训练和测试数据为 SIDD-Medium,需要下载并分 patch.

已将分好 patch 的数据放在了 Ai Studio 里.

可直接运行下面的脚本解压:

In [ ]
!cd data && tar -xf data140841/SIDD_patches.tar.gz
登录后复制    

3.2 预训练模型

MIRNet V1:

最新预训练模型,已转为 paddle 的,名为 MIRNetV1_torch.pdparams.复现的模型,名为 MIRNetV1_paddle.pdparams.pytorch 的初始化参数,名为 torch_init.pdparams

MIRNet V2:

复现的模型,名为 MIRNetV2_paddle.pdparams,导出的静态图模型参数包括 model.pdmodel 和 model.pdiparams.

运行以下脚本解压:

In [ ]
!unzip data/data150163/pretrained_models.zip -d work/pretrained_models
登录后复制    

3.3 文件结构

MIRNet_Paddle    |-- configs                       # 单机单卡/四卡训练配置文件    |-- dataloaders                # 数据集相关文件    |-- SIDD_patches         |-- train                      # SIDD-Medium 训练数据         |-- val                         # SIDD 测试数据    |-- networks         |-- MIRNet_model.py          # MIRNetV1模型代码         |-- MIRNet_V2_model.py    # MIRNetV2模型代码    |-- pretrained_models              # 预训练模型    |-- utils                                     # 一些工具代码    |-- config.py                             # 配置文件    |-- losses.py                             # 损失函数    |-- test_denoising_sidd.py       # 测试SIDD数据上的指标    |-- train_denoising.py              # 训练代码
登录后复制    

4. 环境依赖

PaddlePaddle >= 2.2.0

scikit-image == 0.19.2

In [ ]
!pip install scikit-image natsort yacs
登录后复制    

5. 核心代码

MIRNet 的核心为 MRB 模块,核心代码为:

class MSRB(nn.Layer):    def __init__(self, n_feat, height, width, stride, bias):        super(MSRB, self).__init__()        self.n_feat, self.height, self.width = n_feat, height, width        self.blocks = nn.LayerList([nn.LayerList([DAU(int(n_feat*stride**i))]*width) for i in range(height)])        INDEX = np.arange(0,width, 2)        FEATS = [int((stride**i)*n_feat) for i in range(height)]        SCALE = [2**i for i in range(1,height)]        self.last_up   = nn.LayerDict()        for i in range(1,height):            self.last_up.update({f'{i}': UpSample(int(n_feat*stride**i),2**i,stride)})        self.down = nn.LayerDict()        self.up   = nn.LayerDict()        i=0        SCALE.reverse()        for feat in FEATS:            for scale in SCALE[i:]:                self.down.update({f'{feat}_{scale}': DownSample(feat,scale,stride)})            i+=1        i=0        FEATS.reverse()        for feat in FEATS:            for scale in SCALE[i:]:                                self.up.update({f'{feat}_{scale}': UpSample(feat,scale,stride)})            i+=1        self.conv_out = nn.Conv2D(n_feat, n_feat, kernel_size=3, padding=1, bias_attr=bias)        self.selective_kernel = nn.LayerList([SKFF(n_feat*stride**i, height) for i in range(height)])    def forward(self, x):        inp = x.clone()        #col 1 only        blocks_out = []        for j in range(self.height):            if j==0:                inp = self.blocks[j][0](inp)            else:                inp = self.blocks[j][0](self.down[f'{inp.shape[1]}_{2}'](inp))            blocks_out.append(inp)        #rest of grid        for i in range(1,self.width):            #Mesh            # Replace condition(i%2!=0) with True(Mesh) or False(Plain)            # if i%2!=0:            tmp=[]            for j in range(self.height):                TENSOR = []                nfeats = (2**j)*self.n_feat                for k in range(self.height):                    TENSOR.append(self.select_up_down(blocks_out[k], j, k))                selective_kernel_fusion = self.selective_kernel[j](TENSOR)                tmp.append(selective_kernel_fusion)            #Forward through either mesh or plain            for j in range(self.height):                blocks_out[j] = self.blocks[j][i](tmp[j])        #Sum after grid        out=[]        for k in range(self.height):            out.append(self.select_last_up(blocks_out[k], k))          out = self.selective_kernel[0](out)        out = self.conv_out(out)        out = out + x        return out    def select_up_down(self, tensor, j, k):        if j==k:            return tensor        else:            diff = 2 ** np.abs(j-k)            if j登录后复制        

损失函数采用 Charbonnier Loss,实现较为简单,代码如下:

class CharbonnierLoss(nn.Layer):    """Charbonnier Loss (L1)"""    def __init__(self, eps=1e-3):        super(CharbonnierLoss, self).__init__()        self.eps = eps    def forward(self, x, y):        diff = x - y        # loss = paddle.sum(paddle.sqrt(diff * diff + self.eps))        loss = paddle.mean(paddle.sqrt((diff * diff) + (self.eps*self.eps)))        return loss
登录后复制    

6. 快速开始

配置文件在work/configs下,可修改学习率、batch_size等参数

对于MIRNet V1(论文复现赛93题),单卡运行的配置文件为MIRNet_1cards.yml:

优化相关的设置:batch_size设置为4,epoch设为60,初始学习率设为2e-4.

OPTIM:  BATCH_SIZE: [4]  NUM_EPOCHS: [60]  LR_INITIAL: 2e-4
登录后复制        

对于MIRNet V2(特色模型挑战赛5),单卡运行的配置文件为MIRNetV2_1cards.yml:

采用Progressive Learning,即逐渐增大输入的patch_size,并减小batch_size,此处patch_size从128增加到256,batch_size从8减小到2,各个patch_size对应epoch数为30、15、10、5.

OPTIM:  BATCH_SIZE: [8,6,4,2]  NUM_EPOCHS: [30,15,10,5]  LR_INITIAL: 2e-4  TRAINING:  PATCH_SIZE: [128,160,192,256]  NUM_WORKERS: [4,4,4,4]
登录后复制        

注:学习策略采用Warm up + Cosine Anneal LR,其中Warm up的epoch数为3.

6.1 模型训练

MIRNet V1单卡运行的代码如下:

In [ ]
!cd work && python train_denoising.py --model MIRNet --gpus 1
登录后复制    

同时给出 MIRNet 单机四卡和 MIRNet V2 的训练脚本,为更好的体验 MIRNet 的训练并得到复现结果,请使用脚本任务.

## MIRnet V1# 单机四卡!cd work && python -m paddle.distributed.launch train_denoising.py --model MIRNet --gpus 4## MIRNet V2# 单机单卡!cd work && python train_denoising.py --model MIRNetV2 --gpus 1# 单机四卡!cd work && python -m paddle.distributed.launch train_denoising.py --model MIRNetV2 --gpus 4
登录后复制        

训练过程会将模型参数保存在 ./ckpt/Denoising/model/ 文件夹下.

6.2 日志读取

训练过程会将日志记录保存在 ./ckpt/Denoising/logs/ 文件夹下,例如 MIRNet V2 的日志目录为 ./ckpt/Denoising/logs/MIRNet_V2/

日志是用 VisualDL 工具记录的,可在 CodeLab 左侧的数据模型可视化中,设置 logdir 查看.

6.3 模型评估

在 SIDD 测试数据上作测试,以 MIRNet V1 为例,若想测试 MIRNet V2,只需将 --model MIRNet 改为 MIRNetV2,同时修改权重weights路径.

In [ ]
# MIRNet V1!cd work && python test_denoising_sidd.py --input_dir ../data/SIDD_patches/val --weights ./pretrained_models/MIRNetV1_paddle.pdparams --model MIRNet
登录后复制    

输出如下:

# MIRNet V1PSNR: 39.6872 SSIM: 0.9586# MIRNet V2PSNR: 39.5286SSIM: 0.9578
登录后复制        

达到了验收精度.

6.4 模型预测

在 SIDD 小验证集上作预测,结果存放在 work/results/ 文件夹下,下以 MIRNet V1 为例,对于 MIRNet V2,同上修改weight和model.

In [ ]
# MIRNet V1!cd work && python predict.py --data_path ./SIDD_patches/val_mini/ --save_path results/ --save_images --model_ckpt ./pretrained_models/MIRNetV1_paddle.pdparams --model MIRNet
登录后复制    

6.5 单张图像去噪测试

导入单张图像,测试去噪效果,首先需要在work/test_images里上传一张图片.

In [41]
# 先上传一张图片,import os.path as ospfrom IPython.display import displayfrom PIL import Imageimg_path = 'bird.png' # 改成自己上传的图片名称full_img_path = osp.join(osp.abspath('work/test_images/'), img_path)img = Image.open(full_img_path).convert('RGB')print('以下为上传的图片:')display(img)
登录后复制        
以下为上传的图片:
登录后复制        
登录后复制                

以 MIRNet V1 为例,对于 MIRNet V2,同上修改model_ckpt和model.

需要指定干净图像和噪声图像,可以只给一张噪声图片,也可以只给一张干净图片,也可以都给.

给定一张噪声图片:指定参数noisy_img,直接输出去噪图片.

给定一张干净图片:指定参数clean_img和noisyL,后者为噪声水平,默认为15,输出加噪图片和去噪图片.

给定噪声图片和干净图片:直接输出去噪图片.

In [ ]
# MIRNet V1 仅给定干净图片,噪声水平为15!cd work && python predict_single.py --clean_img $full_img_path --save_images --model_ckpt ./pretrained_models/MIRNetV1_paddle.pdparams --model MIRNet
登录后复制    In [45]
# 去噪效果查看import globfrom IPython.display import displayfrom PIL import Imageimgs = glob.glob('work/test_images/*')for path in imgs:    print(path)    img = Image.open(path)    display(img)
登录后复制        
work/test_images/bird_noised.png
登录后复制        
登录后复制                
work/test_images/bird_denoised.png
登录后复制        
登录后复制                
work/test_images/bird.png
登录后复制        
登录后复制                

热门推荐

更多

热门文章

更多

首页  返回顶部

本站所有软件都由网友上传,如有侵犯您的版权,请发邮件youleyoucom@outlook.com