参考老师发的代码对代码进行初步理解。
以前的图像去噪大多使用合成数据,这篇文章研究了CNN在真实图像上的去噪效果,其主要贡献在于以下几点:
下载数据:
! wget https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/mini_denoise_dataset.zip
! unzip mini_denoise_dataset
引入基本的库:
import os, time, scipy.io, shutil
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import glob
import re
import cv2
我们使用CBDNet网络,包括一个噪声估计子网络和一个非盲去噪子网络,来实现图像的盲去噪。
CBDNet网络由由一个全卷积网络FCN和一个 UNet 组成,FCN为五层全卷积网络用于噪声估计噪声水平图,第二部分为带残差的UnNet,用于降噪。
网络结构如下所示:
首先是定义FCN全卷积网络,FCN网络结构主要分为两个部分:全卷积部分和反卷积部分。其中全卷积部分为一些经典的CNN网络,用于提取特征;反卷积部分则是通过上采样得到原尺寸的语义分割图像,解决卷积和池化导致图像尺寸的变小。包括 5 次 conv 操作,使用 3x3 的卷积核,使用了1个像素的padding来保证尺寸一致,feature map 数量依次为:3 ==> 32 ==> 32 ==> 32 ==> 32 ==> 3。
class FCN(nn.Module):
def __init__(self):
super(FCN, self).__init__()
# 3 ==> 32 的输入卷积
self.inc = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(inplace=True))
# 32 ==> 32 的中间卷积
self.conv = nn.Sequential(
nn.Conv2d(32, 32, 3, padding=1),
nn.ReLU(inplace=True)
)
# 32 ==> 3 的输出卷积
self.outc = nn.Sequential(
nn.Conv2d(32, 3, 3, padding=1),
nn.ReLU(inplace=True)
)
def forward(self, x):
# 第 1 次卷积
conv1 = self.inc(x)
# 第 2 次卷积
conv2 = self.conv(conv1)
# 第 3 次卷积
conv3 = self.conv(conv2)
# 第 4 次卷积
conv4 = self.conv(conv3)
# 第 5 次卷积
conv5 = self.outc(conv4)
return conv5
class single_conv(nn.Module):
def __init__(self, in_ch, out_ch): # 两个参数
super(single_conv, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1), # 二维卷积3×3,1个像素的 # padding保证尺寸
nn.ReLU(inplace=True)) # ReLU函数
def forward(self, x):
x = self.conv(x)
return x
网络分为两次上采样和两次下采样。
下采样:使用2×2的均值pooling
上采样:采用反卷积。
编写up类,需要一个参数in_ch输入通道数,最终输出的通道数是in_ch//2(向下取整除)。
这里的forward函数需要两个参数:
x1:上采样的小尺寸feature map
x2:以前的大尺寸feature map
因为中间的 pooling 可能损失了边缘像素,所以上采样以后的 x1 可能会比 x2 尺寸小。
class up(nn.Module):
def __init__(self, in_ch):
super(up, self).__init__()
self.up = nn.ConvTranspose2d(in_ch, in_ch//2, 2, stride=2) # 反卷积
def forward(self, x1, x2):
# x1 上采样
x1 = self.up(x1)
# 输入数据是四维的,第一个维度是样本数,剩下的三个维度是 CHW
# 所以 Y 方向上的悄寸差别在 [2], X 方向上的尺寸差别在 [3]
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
# 给 x1 进行 padding 操作
x1 = F.pad(x1, (diffX // 2, diffX - diffX//2,
diffY // 2, diffY - diffY//2))
# 把 x2 加到反卷积后的 feature map
x = x2 + x1
return x
这里加入pad是为了避免两个feature map尺寸不一样产生的问题。
编写输出层的类outconv,输出为64个feature map,再用1×1的卷积(起到降维作用,得到降噪效果)变成3个feature map。
最后一层不使用激活函数。
class outconv(nn.Module):
def __init__(self, in_ch, out_ch):
super(outconv, self).__init__()
self.conv = nn.Conv2d(in_ch, out_ch, 1)
def forward(self, x):
x = self.conv(x)
return x
将上述几个类进行整合,形成最终UNet类。
class UNet(nn.Module):
def __init__(self):
super(UNet, self).__init__()
self.inc = nn.Sequential(
single_conv(6, 64),
single_conv(64, 64))
self.down1 = nn.AvgPool2d(2) # 2×2的均值pooling
self.conv1 = nn.Sequential(
single_conv(64, 128),
single_conv(128, 128),
single_conv(128, 128))
self.down2 = nn.AvgPool2d(2) # 2×2的均值pooling
self.conv2 = nn.Sequential(
single_conv(128, 256),
single_conv(256, 256),
single_conv(256, 256),
single_conv(256, 256),
single_conv(256, 256),
single_conv(256, 256))
self.up1 = up(256) # cov2反卷积
self.conv3 = nn.Sequential(
single_conv(128, 128),
single_conv(128, 128),
single_conv(128, 128))
self.up2 = up(128) # cov3反卷积
self.conv4 = nn.Sequential(
single_conv(64, 64),
single_conv(64, 64))
self.outc = outconv(64, 3)
def forward(self, x):
# input conv : 6 ==> 64 ==> 64
inx = self.inc(x)
# 均值 pooling, 然后 conv1 : 64 ==> 128 ==> 128 ==> 128
down1 = self.down1(inx)
conv1 = self.conv1(down1)
# 均值 pooling,然后 conv2 : 128 ==> 256 ==> 256 ==> 256 ==> 256 ==> 256 ==> 256
down2 = self.down2(conv1)
conv2 = self.conv2(down2)
# up1 : conv2 反卷积,和 conv1 的结果相加,输入256,输出128
up1 = self.up1(conv2, conv1)
# conv3 : 128 ==> 128 ==> 128 ==> 128
conv3 = self.conv3(up1)
# up2 : conv3 反卷积,和 input conv 的结果相加,输入128,输出64
up2 = self.up2(conv3, inx)
# conv4 : 64 ==> 64 ==> 64
conv4 = self.conv4(up2)
# output conv: 65 ==> 3,用1x1的卷积降维,得到降噪结果
out = self.outc(conv4)
return out
定义完FCN部分与UNet部分后,整体定义CBDNet网络。
先将数据输入 FCN,得到估计的噪声强度:noise_level,为 3 通道。然后将 3通道的原图像,和 noise_level 拼接在一起,作为 UNet 的输入。经过UNet操作,得到 out ,out 是噪声的 residual mapping,和输入图像加在一起,输出最终的去噪图像。
采用了一 residual learning 的思想,认为噪声的 residual mapping 学习起来更加容易。
class CBDNet(nn.Module):
def __init__(self):
super(CBDNet, self).__init__()
self.fcn = FCN()
self.unet = UNet()
def forward(self, x):
noise_level = self.fcn(x)
concat_img = torch.cat([x, noise_level], dim=1)
out = self.unet(concat_img) + x
return noise_level, out
基于focal loss解决正负样本不平衡问题,focal loss的改进版被提出,这是一种非对称的loss,即Asymmetric Loss。
非对称loss的思想是加大对低估网络噪声等级的惩罚,损失函数定义为:
Asymmetric Loss解决了多标签分类任务中,正负样本不平衡问题,标签错误问题;通过梯度分析,对该loss进行了分析;提出了自适应的方法来控制非对称的级别,简化了超参数选择过程;此外Asymmetric Loss十分高效,容易使用。相比于最近的其他方法,该方法基于主流的网络结构,并且不需要其他的信息。
用在图像上,Total Variation loss可以使图像变得平滑。信号处理中,总变差去噪,也称为总变差正则化,是最常用于数字图像处理的过程,其在噪声去除中具有应用。
它基于这样的原理:具有过多和可能是虚假细节的信号具有高的总变化,即,信号的绝对梯度的积分是高的。根据该原理,减小信号的总变化,使其与原始信号紧密匹配,去除不需要的细节,同时保留诸如边缘的重要细节。其函数定义为:
这种噪声消除技术优于简单的技术,例如线性平滑或中值滤波,这些技术可以降低噪声,但同时可以更大或更小程度地消除边缘。相比之下,即使在低信噪比下,总变差去噪在同时保留边缘同时平滑掉平坦区域中的噪声方面也非常有效。
所谓Reconstruction Loss (重构损失),其实就是CycleGAN中的cycle consistency loss(循环一致性损失)。
上述两种loss并不能保证保留输入图像上与域转换无关的内容不变,只改变与域转换相关的部分,所以在模型中的生成器引入了重构损失。其函数定义为:
Reconstruction loss得到了去噪图像和真实图像之间的差距。
三个部分加在一起,就是最终的损失函数,下面结合代码来看。
class fixed_loss(nn.Module):
def __init__(self):
super().__init__()
def forward(self, out_image, gt_image, est_noise, gt_noise, if_asym):
# 分别得到图像的高度和宽度
h_x = est_noise.size()[2]
w_x = est_noise.size()[3]
# 每个样本为 CHW ,把 H 方向第一行的数据去掉,统计一下一共多少元素
count_h = self._tensor_size(est_noise[:, :, 1:, :])
# 每个样本为 CHW ,把 W 方向第一列的数据去掉,统计一下一共多少元素
count_w = self._tensor_size(est_noise[:, :, : ,1:])
# H 方向,第一行去掉得后的矩阵,减去最后一行去掉后的矩阵,即下方像素减去上方像素,平方,然后求和
h_tv = torch.pow((est_noise[:, :, 1:, :] - est_noise[:, :, :h_x-1, :]), 2).sum()
# W 方向,第一列去掉得后的矩阵,减去最后一列去掉后的矩阵,即右方像素减去左方像素,平方,然后求和
w_tv = torch.pow((est_noise[:, :, :, 1:] - est_noise[:, :, :, :w_x-1]), 2).sum()
# 求平均,得到平均每个像素上的 tvloss
tvloss = h_tv / count_h + w_tv / count_w
loss = torch.mean( \
# 第三部分:重建损失
torch.pow((out_image - gt_image), 2)) + \
# 第一部分:对比损失 第二部分:起平滑作用的 tvloss
if_asym * 0.5 * torch.mean(torch.mul(torch.abs(0.3 - F.relu(gt_noise - est_noise)), torch.pow(est_noise - gt_noise, 2))) + 0.05 * tvloss
return loss
def _tensor_size(self,t):
return t.size()[1]*t.size()[2]*t.size()[3]
从上面的代码中可以看到,对比损失前系数为 0.5, alpha 取值为 0.3,tvloss 系数为 0.05,和论文里的默认参数一致。
对于 gt_noise,只有在使用合成数据进行训练时才会用到;以前的图像去噪,大多在真实图像上加一个随机Gauss噪声,得到噪声图像,这时 gt_noise 是已知的,就能够输入。而这里处理的是真实图像,因此没有 gt_noise,所以在训练时,gt_noise 一直是0。
下面是两个程序中要用到的两个小函数:
# 这个类用于存储 loss,观察结果时使用# 每轮训练一张图像,就计算一下 loss 的均值存储在 self.avg 里,用于输出观察变化# 同时,把当前 loss 的值存储在 self.val 里class AverageMeter(object):
# 初始化
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
#更新
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
# 图像矩阵由 hwc 转换为 chw
def hwc_to_chw(img):
return np.transpose(img, axes=[2, 0, 1])
# 图像矩阵由 chw 转换为 hwc
def chw_to_hwc(img):
return np.transpose(img, axes=[1, 2, 0])
一方面,合成噪声图像单独训练出的模型对真实噪声图像的去噪效果偏弱,而另一方面,真实噪声图像通过大量同场景图片平均得到的无噪声图有些过于光滑,所以将两者结合起来进行网络的训练可以提高网络对于真实噪声图像的泛化能力。因此将真实噪声图像和合成噪声图像一起作为训练集,交替对网络进行训练以提升网络的性能。
初始化基本变量:
# 训练的时候,输入图像尺寸都是 ps x ps 的
ps = 256
train_dir = './mini_denoise_dataset/train/'
train_fns = glob.glob(train_dir + 'Batch_*')
origin_imgs = [None] * len(train_fns)
noised_imgs = [None] * len(train_fns)
定义网络模型、优化器和损失函数:
# 使用GPU训练
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 创建 模型 + 优化器 + 损失函数
model = CBDNet().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = fixed_loss()
开始训练:
cnt = 0
total_loss = AverageMeter()
# 设置为训练模式,即启用 BatchNormalization 和 Dropout
model.train()
for epoch in range(200):
# 内存中清空图片
for i in range(len(train_fns)):
origin_imgs[i] = []
noised_imgs[i] = []
# 打乱训练图片的顺序
for idx in np.random.permutation(len(train_fns)):
# 读入origin image;RGB通道反过来,然后归一化;转化为 float32类型
origin_img = cv2.imread(glob.glob(train_fns[idx] + '/*Reference.bmp')[0])
origin_img = origin_img[:,:,::-1] / 255.0
origin_imgs[idx] = np.array(origin_img).astype('float32')
# 读入noised image;因为一个文件夹里有2张噪声图,这里写了一个循环
train_noised_list = glob.glob(train_fns[idx] + '/*Noisy.bmp')
for nidx in range(len(train_noised_list)):
noised_img = cv2.imread(train_noised_list[nidx])
noised_img = noised_img[:,:,::-1] / 255.0
noised_img = np.array(noised_img).astype('float32')
noised_imgs[idx].append(noised_img)
H, W, C = origin_img.shape
# 从图像中随机取 256x256 大小的块
xx = np.random.randint(0, W-ps+1)
yy = np.random.randint(0, H-ps+1)
temp_origin_img = origin_imgs[idx][yy:yy+ps, xx:xx+ps, :]
temp_noised_img = noised_imgs[idx][nidx][yy:yy+ps, xx:xx+ps, :]
# 生成 0,1 随机数,随机做图像的左右、上下、通道翻转,增加训练样本的多样性
if np.random.randint(0, 2) == 1: # 左右翻转
temp_origin_img = np.flip(temp_origin_img, axis=1)
temp_noised_img = np.flip(temp_noised_img, axis=1)
if np.random.randint(0, 2) == 1: # 上下翻转
temp_origin_img = np.flip(temp_origin_img, axis=0)
temp_noised_img = np.flip(temp_noised_img, axis=0)
if np.random.randint(0, 2) == 1: # 通道翻转
temp_origin_img = np.transpose(temp_origin_img, (1, 0, 2))
temp_noised_img = np.transpose(temp_noised_img, (1, 0, 2))
temp_noised_img_chw = hwc_to_chw(temp_noised_img)
temp_origin_img_chw = hwc_to_chw(temp_origin_img)
cnt += 1
# 这里给输入数据增加一个维度,即原来是三维的,现在是四维的,方便CNN处理
input_var = torch.from_numpy(temp_noised_img_chw.copy()).type(torch.FloatTensor).unsqueeze(0).to(device)
target_var = torch.from_numpy(temp_origin_img_chw.copy()).type(torch.FloatTensor).unsqueeze(0).to(device)
# 噪声图像输入网络处理
noise_level_est, output = model(input_var)
# 计算损失
loss = criterion(output, target_var, noise_level_est, 0, 0)
total_loss.update(loss.item())
# 常规操作: 梯度归零 + 反向传播 + 优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('[Epoch %d] [Img count %d] [Loss.val: %.4f] ([loss.avg: %.4f])\t' % (epoch, cnt, total_loss.val, total_loss.avg))
一些基本参数:
test_dir = './mini_denoise_dataset/test/'
test_fns = glob.glob(test_dir + '*.bmp')
# 建立 result 目录,保存图片处理结果
result_dir = './result/'
if not os.path.exists( result_dir ):
os.mkdir( result_dir )
开始测试:
for ind, test_img_path in enumerate(test_fns):
model.eval()
with torch.no_grad():
print(test_img_path)
# 读入图像,切换RGB通道并归一化,转化为 numpy float32格式
noisy_img = cv2.imread(test_img_path)
noisy_img = noisy_img[:,:,::-1] / 255.0
noisy_img = np.array(noisy_img).astype('float32')
# 转化为 chw 才符合 pytorch 网络的输入格式
temp_noisy_img_chw = hwc_to_chw(noisy_img)
# 图像放到 gpu 上
input_var = torch.from_numpy(temp_noisy_img_chw.copy()).type(torch.FloatTensor).unsqueeze(0).to(device)
# 输入模型得到结果
_, output = model(input_var)
# 输出结果转化为 numpy ,同时,把数据转到 0,1 之间(因为可能会有一些异常值)
output_np = output.squeeze().cpu().detach().numpy()
output_np = chw_to_hwc(np.clip(output_np, 0, 1))
# 把噪声图像,和降噪后的图像拼接在一起,然后保存图像
tempImg = np.concatenate((noisy_img, output_np), axis=1)*255.0
Image.fromarray(np.uint8(tempImg)).save(fp=result_dir + 'test_%d.jpg'%(ind), format='JPEG')
全员对代码进行理解
已经完成的工作:对进行降噪的代码进行全面理解
计划完成的工作:开始前端、后端设计
遇到的困难:将代码全部输入后没有对应图片输出,决定下周研究一下是什么问题