import argparse
from PIL import Image
from skimage import util, img_as_float, io
import os


def add_noise(img_path, img_dir):
    if img_dir == '':
        # 读取图像
        img_org = Image.open(img_path)
        # 转换为 skimage 可操作的格式
        img = img_as_float(img_org)
        image_gaussian = util.random_noise(img, mode="gaussian")  # 加高斯噪声
        image_sp = util.random_noise(img, mode="s&p")  # 加椒盐噪声
        # 存储文件到新的路径中，并修改文件名
        io.imsave(str(img_path)[0:-4] + "_guassian.png", image_gaussian)
        io.imsave(str(img_path)[0:-4] + "_sp.png", image_sp)
    else:
        img_list = os.listdir(img_dir)
        for img_path in img_list:
            img_p = img_dir + '/' + img_path
            # 读取图像
            img_org = Image.open(img_p)
            # 转换为 skimage 可操作的格式
            img = img_as_float(img_org)
            image_gaussian = util.random_noise(img, mode="gaussian")  # 加高斯噪声
            image_sp = util.random_noise(img, mode="s&p")  # 加椒盐噪声

            # 存储文件到新的路径中，并修改文件名
            io.imsave(img_dir + '/' + str(img_path).split('.')[0] + "_guassian.png", image_gaussian)
            io.imsave(img_dir + '/' + str(img_path).split('.')[0] + "_sp.png", image_sp)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_dir', type=str, default='./img')
    parser.add_argument('--img_path', type=str, default='')
    args = parser.parse_args()
    add_noise(args.img_path, args.img_dir)