【CV】基于DnCNN的图像和视频去噪

机器学习初学者

共 5969字,需浏览 12分钟

 ·

2022-04-18 07:40



简介


随着数字图像数量的增加,对高质量的图像需求也在增加。然而,现代相机拍摄的图像会因噪声而退化。图像中的噪声是图像中颜色信息的失真,噪声是指数字失真。当在夜间拍摄时,图像变得更嘈杂。该案例研究试图建立一个预测模型,该模型将带噪图像作为输入并输出去噪后的图像。


深度学习的使用


这个问题是基于计算机视觉的,CNN等深度学习技术的进步已经能够在图像去噪方面提供最先进的性能,用于执行图像去噪的模型是DnCNN(去噪卷积神经网络)。


数据集


BSD300和BSD500数据集均用作训练数据,BSD68用于验证数据。由于数据有限,每个图像使用了4次,即缩放到[1.0,0.7,0.8,0.7]。


每个缩放图像被分割成50x50的块,步幅为20。每个贴片都添加了一个标准偏差在[1,55]之间的高斯噪声。数据生成代码如下所示:

#Fix Noisestddevs = np.random.uniform(1, 55.0, 125000)[:, np.newaxis, np.newaxis, np.newaxis]noise = np.random.normal(loc = 0, scale=stddevs, size=(125000, 50, 50, 3)).astype(np.float16)
def get_dataset(img_path):def image_generator(): patch_size = 50 stride = 20 index = 0for scale in [1, 0.9, 0.8, 0.7]:for path in img_path: true_img = cv2.imread(path)for i in range(0, true_img.shape[0] - patch_size + 1, stride):for j in range(0, true_img.shape[1] - patch_size + 1, stride): Y = true_img[i:i+patch_size, j:j+patch_size] gauss_noise = noise[index].astype(np.float32) X = np.clip(Y + gauss_noise, 0, 255.0) index = (index + 1)%125000yield (X/255.0,),Y/255.0return tf.data.Dataset.from_generator(image_generator, output_signature=((tf.TensorSpec(shape=(None, None, 3)),), (tf.TensorSpec(shape=(None, None, 3)))))

DnCNN体系结构


DnCNN中有三种类型的层:

  1. Conv+ReLU:过滤器大小为3,过滤器数量为64,跨步为1,使用零填充保持卷积后的输出形状,使用ReLU作为激活函数。输出为形状(批量大小,50、50、64)

  2. Conv+批量归一化+ReLU:过滤器大小为3,过滤器数量为64,步长为1,使用零填充保持卷积后的输出形状,使用批量归一化层更好地收敛,ReLU作为激活函数。输出为形状(批次大小,50、50、64)。

  3. Conv:滤镜大小为3,跨步为1,滤镜数量为c(彩色图像为3个,灰度图像为1个),使用零填充在卷积后保持输出形状。输出形状为(批次大小,50,50,c)。

DnCNN模型的输出为残差图像。因此,原始图像=噪声图像-残差图像。


在DnCNN中,在每层卷积之前填充零,以确保中间层的每个特征贴图与输入图像具有相同的大小。根据本文,简单的零填充策略不会导致任何边界伪影。

本文建议深度为17,但本案例研究适用于深度为12和深度为8。


评价指标


评估指标是PSNR(峰值信噪比)分数。它只是一个数值,表示构造的去噪图像与原始图像相比有多好。


模型训练

def get_model(depth, channels):noise_inp = tf.keras.layers.Input(shape = (50, 50, channels), dtype=tf.float32)init = 'Orthogonal'
y = tf.keras.layers.Conv2D(filters = 64, kernel_size = 3, padding = 'same', kernel_initializer=init, use_bias=True)(noise_inp)y = tf.keras.layers.ReLU()(y)for i in range(1, depth-1):y = tf.keras.layers.Conv2D(filters = 64, kernel_size = 3, padding = 'same', kernel_initializer=init, use_bias=True)(y)bn = tf.keras.layers.BatchNormalization(axis=-1, epsilon=1e-5, momentum=0.9)y = bn(y)y = tf.keras.layers.ReLU()(y)residual = tf.keras.layers.Conv2D(filters = channels, kernel_size = 3, padding = 'same', kernel_initializer=init, use_bias=True)(y)
true_img = tf.keras.layers.Subtract()([noise_inp, residual])model = tf.keras.Model(inputs = [noise_inp], outputs=[true_img])model.compile(optimizer=tf.keras.optimizers.Adam(), loss='mse')
return model
def lr_decay(epoch):lr = 1e-3if epoch+1 > 20:lr/=30elif epoch+1 > 10:lr /= 10return lrmodel = get_model(8, 3)lr_callback = tf.keras.callbacks.LearningRateScheduler(lr_decay)dataset = get_dataset(bsd500).shuffle(1000).batch(128).prefetch(tf.data.experimental.AUTOTUNE).repeat(None)model.compile(optimizer=tf.keras.optimizers.Adam(), loss='mse')history = model.fit(x = dataset, steps_per_epoch=2000, epochs=30, shuffle=True,verbose=1,callbacks=[lr_callback])

批量大小=128,每个历元的步数=2000,历元数=30。


结果





BSD68数据集上的峰值信噪比对于标准差25为~28,对于标准差50为~25。

如果深度=12,则BSD68数据集上的峰值信噪比对于标准差25为28.30,对于标准差50为26.13。


应用:视频去噪


我们可以将这个想法扩展到视频帧,每个帧作为输入传递给DnCNN模型,生成的帧传递给视频编写器。

import sysimport tensorflow as tfimport numpy as npimport cv2import timeimport matplotlib.pyplot as pltimport osimport globimport seaborn as snbimport refrom skvideo.io import FFmpegWriter

class Denoiser:def __init__(self, merge_outputs): self.model = tf.keras.models.load_model('./model') self.merge_outputs = merge_outputs
def get_patches(self, frame): patches = np.zeros(shape=(self.batch_size, 50, 50, 3)) counter = 0for i in range(0, self.SCALE_H, 50):for j in range(0, self.SCALE_W, 50): patches[counter] = frame[i:i+50, j:j+50] counter+=1return patches.astype(np.float32)
def reconstruct_from_patches(self, patches, h, w, true_h, true_w, patch_size): img = np.zeros((h,w, patches[0].shape[-1])) counter = 0for i in range(0,h-patch_size+1,patch_size):for j in range(0,w-patch_size+1,patch_size): img[i:i+patch_size, j:j+patch_size, :] = patches[counter] counter+=1return cv2.resize(img, (true_w, true_h), cv2.INTER_CUBIC)
def denoise_video(self, PATH): self.cap = cv2.VideoCapture(PATH) self.H, self.W = int(self.cap.get(4)), int(self.cap.get(3)) self.SCALE_H, self.SCALE_W = (self.H//50 * 50), (self.W//50 * 50) self.batch_size = ((self.SCALE_H * self.SCALE_W) // (50**2))
outputFile = './denoise.mp4' writer = FFmpegWriter( outputFile, outputdict={'-vcodec':'libx264','-crf':'0','-preset':'veryslow' } )
while True: success, img = self.cap.read()if not success:break resize_img = cv2.resize(img, (self.SCALE_W, self.SCALE_H), cv2.INTER_CUBIC).astype(np.float32)
noise_img = resize_img/255.0 patches = self.get_patches(noise_img).astype(np.float32) predictions = np.clip(self.model(patches), 0, 1) pred_img = (self.reconstruct_from_patches(predictions, self.SCALE_H, self.SCALE_W, self.H, self.W, 50)*255.0)if self.merge_outputs: merge = np.vstack([img[:self.H//2,:,:], pred_img[:self.H//2,:,:]]) writer.writeFrame(merge[:,:,::-1])else: writer.writeFrame(pred_img[:,:,::-1]) writer.close()
PATH = sys.argv[1]print(f"Path is : {PATH}")denoise = Denoiser(merge_outputs = True)x = denoise.denoise_video(PATH)

参考

  1. https://arxiv.org/pdf/1608.03981.pdf

  2. https://www.appliedaicourse.com/


GITHUB代码链接:https://github.com/saproovarun/DnCNN-Keras


往期精彩回顾




浏览 48
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报