import os
import os.path
import numpy as np
import skimage
import matplotlib.pyplot as plt

dirname = "./ILSVRC2012_img_val"


save_dirname_noise_v1 = "./training_data/Imagenet_Noisy"
save_dirname_clean = "./training_data/Imagenet_Clean"
filenames = os.listdir(dirname)
filenames.sort()
os.makedirs(save_dirname_noise_v1, exist_ok=True)
os.makedirs(save_dirname_clean, exist_ok=True)

def add_noises(img):
    #Poisson
    poisson_noised = 255.0 * np.random.poisson(30.0*(img/255.0))/30.0
    #Gaussian
    gussian_noised = poisson_noised + np.random.normal(0, 60, img.shape)
    #Bernoulli
    bernoulli_noise_map = np.random.binomial(1, 0.5, img.shape)*255
    bernoulli_noised = np.where(np.random.uniform(0, 1, img.shape)<0.2, bernoulli_noise_map, gussian_noised)
    #Clamp and Quantile
    noised_img = np.clip(bernoulli_noised, 0, 255.0)
    return noised_img.astype('uint8')

for file_name in filenames[:20000]:
    dir_name = os.path.join(dirname,file_name)
    save_dirname_noise_v1_file = os.path.join(save_dirname_noise_v1,file_name)[:-4]+"npy"
    save_dirname_clean_file = os.path.join(save_dirname_clean,file_name)[:-4]+"npy"
    img = plt.imread(dir_name)
    noisy_v1 = add_noises(img.copy()).copy()
    np.save(save_dirname_clean_file,img)
    np.save(save_dirname_noise_v1_file,noisy_v1)
