图像去噪——PMRID训练自己数据集及推理测试(详细图文教程)

一、源码包准备

RMRID最原始的版本是MegEngine框架的,后有了Pytorch版本,本教程主要是Pytorch版本的。

官网MegEngine版本源码:PMRID(MegEngine)

官网Pytorch版本源码:PMRID(Pytorch)

论文地址:论文

我自己提供的源码包是在官网基础上修改过的,其中有一些我添加的代码,我提供的源码包链接为:网盘源码包 ,提取码:he6m

我提供的源码包下载并解压后的文件内容如下:
在这里插入图片描述

二、数据集准备

数据集有很多,看自己需要哪个数据集,官网提到的Kaggle数据集下载链接为:Kaggle

打开网址后直接下载即可,如下:

在这里插入图片描述
如果你想用.RAW格式的原始数据训练,在读入数据这一部分,得修改代码,读入.RAW格式数据的代码如下,关于读取.RAW格式数据和.dng格式数据的方法,查看另外一篇博文读取.RAW数据和.dng数据

在这里插入图片描述

如果数据集是.jpg,.png,.bmp等格式的8bit数据,则保持默认即可,如下:

在这里插入图片描述

2.1 提取数据集名称

在提供的脚本文件中,通过generate_list_sidd.py,将含噪图像名称和真实无噪图像名称提取到一个.txt文件中,如下:

在这里插入图片描述

从生成好的train_list_sidd.txt文件中随机剪切一部分到其它两个文件Sony_val_list.txt和Sony_test_list.txt中,作为验证集和测试集,如下:

在这里插入图片描述

2.2 .txt报错问题

上一步骤中提取图片名称到.txt文件中,注意.txt文件中最后一行的空行一定要删除,不然就会报错,如下:

2.2.1 正确格式

在这里插入图片描述

2.2.2 错误格式

在这里插入图片描述

三、修改配置参数

可以修改的超参数如下

训练轮数:epoch
训练批次大小:batch_size
学习率:learning_rate
使用GPU或CPU:device
日志文件路径:logs_path
模型保存路径:params_path
训练集路径:train_list_path
验证集路径:value_list_path
是否加载预训练权重文件:is_load_pretrained
预训练训中路径:pretrained_path

源码包中具体修改的地方如下,在mian.py脚本中:

在这里插入图片描述

四、训练及保存模型权重

4.1 训练

修改好上面参数后就可以开始训练了,训练前先打开训练命令,如下:

在这里插入图片描述

运行脚本main.py后就开始训练了,如下:

在这里插入图片描述

4.2 保存模型权重文件

训练过程中会将每个epoch的模型权重文件保存在params路径下,如下:

在这里插入图片描述
模型大小为3.98M

五、模型推理测试

训练好后,进行推理前先打开测试命令,如下:

在这里插入图片描述

5.1 导入测试集

由于代码架构设计,验证集的对应的变量名和测试集对应的变量名是一个,所以直接将测试集路径替换掉验证集路径,如下:

在这里插入图片描述

5.2 测试

修改好上面参数后,运行脚本main.py,会在output路径中生成测试结果:

在这里插入图片描述

5.3 测试结果

5.3.1 测试场景1

左侧为原始含噪图像,右侧为模型测试结果

在这里插入图片描述

5.3.2 测试场景2

在这里插入图片描述

5.4 推理速度

5.4.1 CPU推理

我测试的电脑为i7-12700H 2.3GHz,测试图片大小为256*256,CPU上测试的平均推理时间为40ms,部分如下:

在这里插入图片描述

5.4.2 GPU推理

GPU上的平均推理时间为:6ms,部分如下:

在这里插入图片描述

六、总结

以上就是使用PMRID训练自己数据集及推理测试的详细图文教程,由于推理速度达不到要求,后续我没有再继续研究,且本方法要结合相机的iso参数生成噪声训练集才能达到效果最好,不同学者的训练测试结果有所不同是正常的。

后续有研究更深入的学者欢迎一起讨论学习。

总结不易,多多支持,谢谢!

猜你喜欢

转载自blog.csdn.net/qq_40280673/article/details/134965249