对YUV序列添加Gauss噪声

对YUV序列加高斯噪声

目的

对YUV序列的Y分量进行加噪。

代码

import os
from glob import glob
import numpy as np
import argparse
from skimage import img_as_float32, img_as_ubyte


parser = argparse.ArgumentParser(prog='Add Gauss Niose to dec sequence Y')
parser.add_argument('--data_dir', default=None, type=str, metavar='PATH', help="path to save the data (default: None)")
args = parser.parse_args(['--data_dir', r'E:\sequence']) #输入序列存储路径

path_all_dec = glob(os.path.join(args.data_dir, '**/BasketballDrive65Frames_256x256_50_21.yuv'), recursive=True) #待处理的序列列表，可进行批处理
path_all_dec = sorted(path_all_dec)  # path_all_noisy中保存升序排序的待加噪文件

print('Number of sequences: {:d}'.format(len(path_all_dec))) # 待处理文件数目
outpath = r'E:\sequence' # 输出路径

for ii in range(len(path_all_dec)):
    w = 256 #帧宽
    h = 256 #帧高
    f_num = 32 #处理帧数
    var = 0.01 #噪声方差，控制加高斯噪声强度
    mean = 0
    y_size = w * h
    uv_size = w * h >> 1
    fp1 = open(path_all_dec[ii], 'rb')
    name = (path_all_dec[ii].split('.yuv')[0]).split('sequence\\')[1] #提取序列名称
    outdir = os.path.join(outpath, name)
    outname = outdir+'_Gauss_0.0008'+'.yuv' #输出序列名称
    fp2 = open(outname, 'wb')
    for num in range(f_num):
        y_int8 = np.zeros(shape=y_size, dtype='uint8', order='C')
        for i1 in range(y_size):
            y_int8[i1] = ord(fp1.read(1))
        im_y_int8 = y_int8.reshape(h, w, 1).astype('uint8')

        uv = np.zeros(shape=uv_size, dtype='uint8', order='C')
        for j1 in range(uv_size):
            uv[j1] = ord(fp1.read(1))
        im_y = img_as_float32(im_y_int8)
        gauss_noise = np.random.normal(mean, var**0.5, im_y.shape) #生成高斯噪声
        im_gauss = im_y + gauss_noise #进行高斯噪声叠加
        im_gauss = np.clip(im_gauss, 0.0, 1.0) #值裁剪
        im_guass_int8 = img_as_ubyte(im_gauss)
        fp2.write(im_guass_int8)
        fp2.write(uv)
        print('finish add guass noise {:d}/{:d}'.format(num,f_num))
    fp2.close()
    fp1.close()

效果

加噪前序列：

加噪后序列：