第八组 Alpha(1/3)
Posted 33鹅
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了第八组 Alpha(1/3)相关的知识,希望对你有一定的参考价值。
第八组 Alpha(1/3)
文章目录
一、完成任务
参考老师发的代码对代码进行初步理解。
1.图像去噪网络 CBDNet
以前的图像去噪大多使用合成数据,这篇文章研究了CNN在真实图像上的去噪效果,其主要贡献在于以下几点:
- 提出了一个更加真实的噪声模型,其考虑了信号依赖噪声和ISP流程对噪声的影响,展示了图像噪声模型在真实噪声图像中起着关键作用。
- 提出了CBDNet模型,其包括了一个噪声估计子网络和一个非盲去噪子网络,可以实现图像的盲去噪(即未知噪声水平)。
- 提出了非对称学习(asymmetric learning)的损失函数,并允许用户交互式调整去噪结果,增强了去噪结果的鲁棒性。
- 将合成噪声图像与真实噪声图像一起用于网络的训练,提升网络的去噪效果和泛化能力。
下载数据:
! 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
2. 定义 CBDNet 网络
我们使用CBDNet网络,包括一个噪声估计子网络和一个非盲去噪子网络,来实现图像的盲去噪。
CBDNet网络由由一个全卷积网络FCN和一个 UNet 组成,FCN为五层全卷积网络用于噪声估计噪声水平图,第二部分为带残差的UnNet,用于降噪。
网络结构如下所示:
2.1 FCN部分
首先是定义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
2.2 UNet部分
- 编写single_conv类用于卷积操作(需要两个参数:in_ch输入通道数,out_ch输出通道数),类中包括卷积和ReLU函数。
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
2.3 CBDNet 分析
定义完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
2.4设计损失函数
第一部分:Asymmetric Loss
基于focal loss解决正负样本不平衡问题,focal loss的改进版被提出,这是一种非对称的loss,即Asymmetric Loss。
非对称loss的思想是加大对低估网络噪声等级的惩罚,损失函数定义为:
Asymmetric Loss解决了多标签分类任务中,正负样本不平衡问题,标签错误问题;通过梯度分析,对该loss进行了分析;提出了自适应的方法来控制非对称的级别,简化了超参数选择过程;此外Asymmetric Loss十分高效,容易使用。相比于最近的其他方法,该方法基于主流的网络结构,并且不需要其他的信息。
第二部分:Total Variation loss
用在图像上,Total Variation loss可以使图像变得平滑。信号处理中,总变差去噪,也称为总变差正则化,是最常用于数字图像处理的过程,其在噪声去除中具有应用。
它基于这样的原理:具有过多和可能是虚假细节的信号具有高的总变化,即,信号的绝对梯度的积分是高的。根据该原理,减小信号的总变化,使其与原始信号紧密匹配,去除不需要的细节,同时保留诸如边缘的重要细节。其函数定义为:
这种噪声消除技术优于简单的技术,例如线性平滑或中值滤波,这些技术可以降低噪声,但同时可以更大或更小程度地消除边缘。相比之下,即使在低信噪比下,总变差去噪在同时保留边缘同时平滑掉平坦区域中的噪声方面也非常有效。
第三部分:Reconstruction 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])
3.初始化基本变量,开始训练
一方面,合成噪声图像单独训练出的模型对真实噪声图像的去噪效果偏弱,而另一方面,真实噪声图像通过大量同场景图片平均得到的无噪声图有些过于光滑,所以将两者结合起来进行网络的训练可以提高网络对于真实噪声图像的泛化能力。因此将真实噪声图像和合成噪声图像一起作为训练集,交替对网络进行训练以提升网络的性能。
初始化基本变量:
# 训练的时候,输入图像尺寸都是 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以上是关于第八组 Alpha(1/3)的主要内容,如果未能解决你的问题,请参考以下文章