Add various noise types to image when using tf.data.dataset

Question

I have been using the function mentioned here to add different types of noise (Gauss, salt and pepper, etc) to an image.

However, I am trying to build an input pipeline using tf.data.Dataset. I think I have figured out how to add Gaussian and Poisson noise:

@tf.function
def load_images(imagePath):   
   label = tf.io.read_file(base_path + "/clean/" + imagePath)   
   label = tf.image.decode_jpeg(label, channels=3)   
   label = tf.image.convert_image_dtype(label, dtype=tf.float32)

   image=label+tf.random.normal(shape=tf.shape(label),mean=0,stddev=0.1**0.5) #Add Gauss noise
   #image=label+tf.random.poisson(shape=tf.shape(label),lam=0.3) #Add Poisson noise

   return (image, label)

How could I add salt & pepper and speckle noise, as well?

Answer 1

Here's a way to do the random salt and pepper augmentation:

from sklearn.datasets import load_sample_image
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

china = load_sample_image('china.jpg')[None, ...].astype(np.float32) / 255

def salt_and_pepper(image, prob_salt=0.1, prob_pepper=0.1):
    random_values = tf.random.uniform(shape=image[0, ..., -1:].shape)
    image = tf.where(random_values < prob_salt, 1., image)
    image = tf.where(1 - random_values < prob_pepper, 0., image)
    return image

ds = tf.data.Dataset.from_tensor_slices(china).batch(1).map(salt_and_pepper)

result_image = next(iter(ds))

plt.imshow(result_image[0])
plt.show()