图片添加高斯噪声和椒盐噪声python

使用Python给图片添加高斯噪声和椒盐噪声，在研究图像降噪算法时，经常会使用到，简单的写了几行代码。

import cv2
import os
import numpy as np

def Expand2Dto3D(img):
    img = np.expand_dims(img, axis=2)
    img = np.repeat(img, 3, axis=2)
    return img

def AddSaltAndPepperNosie(img, pro):
    noise = np.random.uniform(0, 255, img[:, :, 0].shape)
    mask = noise < pro * 255
    mask = Expand2Dto3D(mask)
    img = img * (1 - mask)

    mask = noise > 255 - pro * 255
    mask = Expand2Dto3D(mask)
    img = 255 * mask + img * (1 - mask)
    return img

def AddGaussNoise(img, sigma, mean=0):
    # 大概率abs(noise) < 3 * sigma
    noise = np.random.normal(mean, sigma, img.shape)
    img = img.astype(np.float)
    img = img + noise
    img = np.clip(img, 0, 255)
    img = img.astype(np.uint8)
    return img

def AddGaussNoiseGray(img, sigma, mean=0):
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)
    noise = np.random.normal(mean, sigma, lab[:, :, 0].shape)
    lab = lab.astype(np.float)
    lab[:, :, 0] = lab[:, :, 0] + noise
    lab[:, :, 0] = np.clip(lab[:, :, 0], 0, 255)
    lab = lab.astype(np.uint8)
    img = cv2.cvtColor(lab, cv2.COLOR_Lab2BGR)
    return img

if __name__ == __main__:
    img = cv2.imread(test3.jpg, 1)
    print(img.shape)
    noiseImg = AddGaussNoise(img, 20, 0)
    cv2.imwrite(test_gauss_noise_color.jpg, noiseImg)

    noiseImgGray = AddGaussNoiseGray(img, 20, 0)
    cv2.imwrite(test_gauss_noise_gray.jpg, noiseImgGray)

    noiseImgSalt = AddSaltAndPepperNosie(img, 0.1)
    cv2.imwrite(test_salt_noise.jpg, noiseImgSalt)

效果如下：

原图

高斯彩噪，sigma=20

仅亮度分量高斯噪声，sigma=20

椒盐噪声，概率10%