图像去噪——CBDNet网络训练自己数据集及推理测试，模型转ONNX模型（详细图文教程）

CBDNet 主要由两个子网络组成：噪声估计子网络和去噪子网络。噪声估计子网络用于估计图像的噪声水平，而去噪子网络用于去除图像中的噪声。

CBDNet 的优势在于：

它采用了更真实的噪声模型，既考虑了泊松-高斯模型，还考虑了信号依赖噪声和 ISP 对噪声的影响。

它采用了非对称损失函数，可以提高网络的泛化能力。

它结合了合成噪声图像和真实噪声图像进行训练，可以更好地适应真实场景。

CBDNet 的劣势在于：

它需要大量的训练数据，训练过程比较耗时。

它对硬件资源要求比较高。

一、源码包准备

官网提供了源码包，我自己也提供了一份，我在官网基础上修改了一些代码，建议学者使用我提供的源码包。本教程是Pytorch版本的。

官网链接：CNDNet

我提供的源码包：网盘，提取码为：7nlv

论文地址：论文

下载解压后的样子如下：

二、环境准备

我自己的训练和测试环境如下，供参考，其它版本也行。

三、数据集准备

3.1 官网数据集

官网教程中有两个数据集，SIDD和Syn，且在链接中提供了一个已经训练好的模型权重文件。数据集和模型权重的下载链接为：SIDD Syn如下：

上面官网提供的两个数据集中，其中SIDD是真实的噪声数据集，Syn是合成噪声数据集。

下载后解压，其中SIDD数据集内容如下：

Syn数据集解压后的样子如下：

3.2 自己数据集准备

官网提供的数据集是将一副高分辨率图像裁剪为256*256大小后再加噪声，每一张高分辨率图像裁剪后得到的小图构成一个子文件夹，多个子文件夹构成整个数据集。

自己制作数据集时可以不用这么小的图片，也不用分这么多子文件夹，只需要将无噪图像和噪声图像同时放到同一个文件夹中，还要注意图片名字命名有规则，命名不一定要按照我的命名方式，自定义规则即可。如下：

如果自己修改数据集的名字，在代码中也要对应修改，不然运行代码找不到图片，代码中具体修改的地方如下：

四、训练

4.1 参数修改

下面是超参数修改：

4.2 训练集路径读取

下面是数据集路径读取：

读入数据这里，batch_size的设置一定要小于子文件夹个数，不然训练时损失函数一直为0，训练结果不对。因为官网提供的源码中，如果一次喂入图片数量小于batch_size，就会被舍弃（主要是因为drop_last=True参数的设置），舍弃了就没数据训练了。

官网读取数据这部分原理是，假设选择SIDD数据集，batch_size设置为64，SIDD中有320个子文件夹，那么一个epoch中，就会随机选取64个子文件，并从64个子文件夹中选取一张图片，迭代5次后完成一个epoch。读取数据的代码如下，这部分代码在./dataset/loader.py脚本中第36行有：

class Real(Dataset):                                                             # 该类继承自Dataset类。Real类用于处理图像数据集，特别是用于处理含有噪声的图像和对应的干净图像
	def __init__(self, root_dir, sample_num, patch_size=128):                    # 定义了类的初始化函数，接受三个参数：root_dir（数据集的根目录），sample_num（样本数量），patch_size（图像块的大小，默认为128）。
		self.patch_size = patch_size                                             # 将传入的patch_size赋值给类的成员变量self.patch_size

		folders = glob.glob(root_dir + '/*')                                     # 获取root_dir目录下的所有文件夹
		folders.sort()                                                           # 对获取到的文件夹进行排序

		self.clean_fns = [None] * sample_num                                     # 初始化一个长度为sample_num的列表self.clean_fns，所有元素都为none
		for i in range(sample_num):                                              # 对于每一个样本
			self.clean_fns[i] = []                                               #  将self.clean_fns的第i个元素设置为一个空列表

		for ind, folder in enumerate(folders):                                   # 对于每一个文件夹
			clean_imgs = glob.glob(folder + '/*GT_SRGB*')
			# clean_imgs = glob.glob(folder + '/*gt_*')  # 获取该文件夹下所有名字中包含GT_SRGB的文件，这些文件是干净的图像。
			clean_imgs.sort()                                                    # 对获取到的干净图像进行排序

			for clean_img in clean_imgs:                                         # 对于每一个干净的图像
				self.clean_fns[ind % sample_num].append(clean_img)               # 将该图像的文件名添加到self.clean_fns的相应列表中

	def __len__(self):                                                           # 定义了类的__len__函数，该函数返回数据集的大小。
		l = len(self.clean_fns)                                                  # 计算self.clean_fns的长度，即数据集的大小
		return l                                                                 # 返回数据集的大小

	def __getitem__(self, idx):                                                  # 定义了类的__getitem__函数，该函数用于获取数据集的第idx个样本。
		clean_fn = random.choice(self.clean_fns[idx])                            # 中随机选择一个干净的图像

		clean_img = read_img(clean_fn)                                           # 读取该干净的图像。
		noise_img = read_img(clean_fn.replace('GT_SRGB', 'NOISY_SRGB'))          # 读取对应的含有噪声的图像
		# noise_img = read_img(clean_fn.replace('gt_', 'noise_'))  # 读取对应的含有噪声的图像

		if self.patch_size > 0:                                                  # 如果patch_size大于0
			[clean_img, noise_img] = get_patch([clean_img, noise_img], self.patch_size)    # 则从干净的图像和含有噪声的图像中获取一个大小为patch_size的图像块

		return hwc_to_chw(noise_img), hwc_to_chw(clean_img), np.zeros((3, self.patch_size, self.patch_size)), np.zeros((3, self.patch_size, self.patch_size))    #  返回含有噪声的图像块、干净的图像块以及两个全零的占位符

官网提供的这种读取方法对于大数据集读取速度很快，如果是自己制作的小数据集，只有一个文件夹时，就没必要按照官网的方法读取，因为如果只有一个文件夹，那batch_size就只能设置为1，batch_size太小，不利于训练和模型的泛化能力。修改读入数据这部分的代码，直接将所有图片数据读取存到列表中，在遍历列表读取即可，这么做的缺点是，数据集较庞大时，前期加载数据集会比较慢。修改读数据代码如下：

实际代码为：

class Real(Dataset):
	def __init__(self, root_dir, sample_num, patch_size=128):
		self.patch_size = patch_size
		self.count = 0
		folders = glob.glob(root_dir + '/*')
		folders.sort()

		# self.clean_fns = [None] * sample_num
		# for i in range(sample_num):
			# self.clean_fns[i] = []
		self.clean_fns = []

		for ind, folder in enumerate(folders):
			# clean_imgs = glob.glob(folder + '/*GT_SRGB*')
			clean_imgs = glob.glob(folder + '/*gt_*')
			clean_imgs.sort()

			for clean_img in clean_imgs:
				# self.clean_fns[ind % sample_num].append(clean_img)
				self.clean_fns.append(clean_img)
			

	def __len__(self):
		l = len(self.clean_fns)
		return l

	def __getitem__(self, idx):                                        # 定义了一个名为 __getitem__ 的方法，它通常用于实现自定义 Python 对象的索引行为，例如访问自定义数据集中的元素。它接受一个索引 idx 作为输入，表示要获取的元素
		# clean_fn = random.choice(self.clean_fns[idx])                       # 从指定索引 idx 处的干净图像文件路径列表中随机选择一个文件路径   干净文件路径列表存储在类的 self.clean_fns 属性中
		clean_fn = random.choice(self.clean_fns)
		# clean_img = read_img(clean_fn)               
		clean_img = read_img(clean_fn)                                      # 使用名为 read_img 的函数（通常用于读取图像）从选定的文件路径读取干净图像
		# noise_img_name = clean_fn.replace('gt_', 'noise_')
		noise_img = read_img(clean_fn.replace('gt_', 'noise_'))     # 通过将干净图像文件路径中的 "GT_SRGB" 替换为 "NOISY_SRGB" 来读取对应的噪声图像

		if self.patch_size > 0:                                             # 检查是否指定了 patch 大小
			[clean_img, noise_img] = get_patch([clean_img, noise_img], self.patch_size)   # 如果指定了 patch 大小，则使用名为 get_patch 的函数从干净图像和噪声图像中提取指定大小的 patch

		return hwc_to_chw(noise_img), hwc_to_chw(clean_img), np.zeros((3, self.patch_size, self.patch_size)), np.zeros((3, self.patch_size, self.patch_size))

用我提供的读数据方法，就可以根据自己电脑性能，设置为较大的batch_size值进行训练。

4.3 单卡或多卡训练

官网提供的源码，默认是直接调用电脑端的所有显卡并行训练，如果想自定义在第二块单卡上训练，需要添加代码，如下：

如果想要多卡并行训练，保持默认即可，不用修改。

4.4 训练

上面参数和路径都修改好后直接运行train.py脚本就开始训练了，如下：

4.5 保存模型权重

训练过程中的模型权重文件会自动保存到根目录下的save_model文件夹中，如下：

五、推理测试

5.1 单帧测试

5.1.1 命令方式

如果只测试一张图片，在终端中输入下面命令：

python predict.py input_filename output_filename

其中input_filename是包含路径的图片名，output_filename是包含保存路径的图片名。实际例子命令如下：

python predict.py Test_Image/ETDS_GaoDe_X4_bmp/4_ETDS_M7C48_x4.bmp Result_image/whq/4_Train_ETDS_M7C48_x4_Denoise.bmp

5.1.2 参数配置方式

如果开发编译环境使用的是Pycharm，也可以使用Configuration参数配置方式测试，如下：

5.2 多帧遍历文件夹测试

如果想直接批量测试一个文件夹中的多张图片，运行我提供的脚本，其中predict_for_CPU.py是CPU批量处理的脚本，predict_for_GPU.py是GPU批量处理的脚本。分为两个脚本是为了下一步的推理时间测试。

批量测试的脚本具体使用如下：

5.3 推理速度

5.3.1 GPU

GPU测试环境：Nvidia GeForce RTX 3050，测试图片96*96，推理时间：2.8ms/fps

5.3.2 CPU

测试环境：12th Gen Intel? Core? i7-12700H 2.30 GHz，测试图片96*96，推理速度：43.61ms/fps

六、转ONNX

为方便部署，将上面训练好的模型权重文件转为ONNX中间格式。

6.1 转换代码

import torch
import torch.nn as nn
import onnx
import numpy as np
from onnx import load_model, save_model
from onnx.shape_inference import infer_shapes
# from models_DnCNN import DnCNN
from model.cbdnet import Network

# 加载模型
# dncnn_model = DnCNN(input_chnl=1, groups=1)
# dncnn_model = torch.load("./model_DnCNN_datav1-sigma11/model_DnCNN_datav1_epoch_500.pth")["model"]
# dncnn_model.load_state_dict(torch.load("./model_DnCNN_datav1/model_DnCNN_datav1_best.pth", map_location="cuda:1")["model"].state_dict())

model = Network()
# state_dict = torch.load('save_model/checkpoint.pth.tar', map_location=torch.device('cpu'))["state_dict"]
state_dict = torch.load('save_model/checkpoint.pth.tar')["state_dict"]

# 创建一个新的state_dict，其键没有'module.'前缀
from collections import OrderedDict

new_state_dict = OrderedDict()

for k, v in state_dict.items():
    name = k[7:]  # 删除'module.'前缀
    print("name", name)
    new_state_dict[name] = v

# 加载新的state_dict
model.load_state_dict(new_state_dict)

# 设置为eval模式，固定bn等操作
# dncnn_model.eval()
# dncnn_model.to("cuda:1")

model.eval()
model.to("cuda:0")
torch.no_grad()
# 设置模型的输入
input = torch.randn((1, 3, 480, 360), dtype=torch.float).to("cuda:0")
torch.onnx.export(model, input, "./Export_ONNX_Result/CDBNet_2.onnx", input_names=["inputs"], output_names=["outputs"], opset_version=14, verbose=1)

# torch.onnx.export(model, input, "./dncnn-sigma11-light.onnx", input_names=["inputs-jl"], output_names=["outputs-jl"], opset_version=14, verbose=1,
#                   dynamic_axes={"inputs-jl":{2:"inputs_height", 3:"inputs_weight"}, "outputs-jl":{2:"outputs_height", 3:"outputs_weight"}})

print("Model has benn converted to onnx")

# onnx_model = load_model("./dncnn-sigma11.onnx")
# onnx_model = infer_shapes(onnx_model)

# save_model(onnx_model, "dncnn-sigma11-shape.onnx")

转后得到的中间格式如下：

6.2 可视化网络结构

可视化网络结构，使用Netron，网址：Netron
打开网络结构如下：

6.3 检验转换后的ONNX模型是否正确

输入同样的数据到转换后的ONNX模型中和原始训练好的模型中，比较两模型的输出差值大不大，在接受范围内就说明转换成功。使用方法及检验代码如下：

检验代码为：

import torch
import numpy as np
import onnxruntime
# from models.ecbsr import ECBSR
# from models.plainsr import PlainSR
from model.cbdnet import Network

def torch_model():
    device = torch.device('cpu')
    ## definitions of model, loss, and optimizer
    # model_ecbsr = ECBSR(module_nums=4, channel_nums=16, with_idt=0, act_type='prelu', scale=4, colors=1).to(device)
    # model_plain = PlainSR(module_nums=4, channel_nums=16, act_type='prelu', scale=4, colors=1).to(device)
    
    # print("load pretrained model: {}!".format("/home/jl/Project/ECBSR/experiments/Visible-light-1channel-noise5-psnr/models/model_x4_514.pt"))
    # model_ecbsr.load_state_dict(torch.load("/home/jl/Project/ECBSR/experiments/Visible-light-1channel-noise5-psnr/models/model_x4_514.pt", map_location='cpu'))

    model = Network()

    print("load pretrained model: {}!".format("save_model/checkpoint.pth.tar"))

    state_dict = torch.load('save_model/checkpoint.pth.tar')["state_dict"]

    # 创建一个新的state_dict，其键没有'module.'前缀
    from collections import OrderedDict

    new_state_dict = OrderedDict()

    for k, v in state_dict.items():
        name = k[7:]  # 删除'module.'前缀
        new_state_dict[name] = v

    # 加载新的state_dict
    model.load_state_dict(new_state_dict)
    return model



    ## copy weights from ecbsr to plainsr
    # depth = len(model_ecbsr.backbone)
    # for d in range(depth):
    #     module = model_ecbsr.backbone[d]
    #     act_type = module.act_type
    #     RK, RB = module.rep_params()
    #     model_plain.backbone[d].conv3x3.weight.data = RK
    #     model_plain.backbone[d].conv3x3.bias.data = RB
    #
    #     if act_type == 'relu':     pass
    #     elif act_type == 'linear': pass
    #     elif act_type == 'prelu':  model_plain.backbone[d].act.weight.data = module.act.weight.data
    #     else: raise ValueError('invalid type of activation!')
    # return model_ecbsr

def pytorch_out(input):
    model = torch_model() #model.eval
    # input = input.cuda()
    # model.cuda()
    torch.no_grad()
    model.eval()
    output = model(input)
    # print output[0].flatten()[70:80]
    out1 = output[0]
    out2 = output[1]
    out = torch.stack((out1, out2))
    return out

def pytorch_onnx_test():
    def to_numpy(tensor):
        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

    # 测试数据
    torch.manual_seed(66)
    dummy_input = torch.randn(1, 3, 480, 360, device='cpu')

    sess = onnxruntime.InferenceSession("./Export_ONNX_Result/CDBNet_2.onnx")

    # onnx 网络输出
    onnx_out = np.array(sess.run(None, {"inputs": to_numpy(dummy_input)}))  #fc 输出是三维列表
    print("==============>")
    print(onnx_out)
    print(onnx_out.shape)
    print("==============>")
    torch_out_res = pytorch_out(dummy_input).detach().numpy()   #fc输出是二维 列表
    print(torch_out_res)
    print(torch_out_res.shape)

    print("===================================>")
    print("输出结果验证小数点后四位是否正确,都变成一维np")

    torch_out_res = torch_out_res.flatten()
    onnx_out = onnx_out.flatten()

    pytor = np.array(torch_out_res,dtype="float32") #need to float32
    onn=np.array(onnx_out,dtype="float32")  ##need to float32
    np.testing.assert_almost_equal(pytor,onn, decimal=5)  #精确到小数点后4位，验证是否正确，不正确会自动打印信息
    print("恭喜你 ^^ , onnx 和 pytorch 结果一致, Exported model has been executed decimal=5 and the result looks good!")

pytorch_onnx_test()

运行上面代码后，输出如下，则说明ONNX模型转换成功，可以直接放到其它平台部署了。

七、测试结果

7.1 测试场景1

7.2 测试场景2

7.3 测试场景3

7.4 测试场景4

八、总结

以上就是图像去噪CBDNet网络训练自己数据集及推理测试，并将训练好的模型转ONNX模型的详细实现过程。网络架构需要花些时间解读，学者仔细研究。

总结不易，多多支持，谢谢！