Pepper & Salt noise image example

pepper_salt_noise.py

#!/usr/bin/env python3
from __future__ import print_function
import sys
try:
    import cv2
    import numpy as np
except ImportError:
    print("install cv2 and numpy please.", file=sys.stderr)
    sys.exit(1)

def pepper_salt(img, p=0.1, value=255):
    index = np.random.binomial(1, p, img.shape[:2])
    noise_img = img.copy()

    if len(img.shape) == 2:
        # gray scale image
        noise_img[np.where(index > 0)] = value
    else:
        # color image
        for channel in range(img.shape[-1]):
            noise_img[:,:,channel][np.where(index > 0)] = value
    return noise_img.astype(np.uint8)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("usage: {} image.png".format(sys.argv[0]))
        sys.exit(1)
    img = cv2.imread(sys.argv[1], cv2.IMREAD_COLOR)

    pepper_salt_img = pepper_salt(img, 0.005)
    cv2.imshow("noise image", pepper_salt_img)
    cv2.waitKey(0)
    cv2.imwrite("noise_image.png", pepper_salt_img)
    sys.exit(0)