基于DnCNN的圖像和視頻去噪

#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)))))

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])

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)