添加salt噪音函数优化

import cv2
import numpy as np
from matplotlib import pyplot as plt

def salt(img, n):
    ilist = np.random.randint(0, img.shape[1],n)
    jlist = np.random.randint(0, img.shape[0],n)
    for k in range(n):
        i = ilist[k]
        j = jlist[k]
        if img.ndim == 2:
           img[j,i] = 255
        elif img.ndim == 3:
            img[j,i,0]= 255
            img[j,i,1]= 255
            img[j,i,2]= 255
    return img


if __name__ == '__main__':
    img = cv2.imread("data/girl.jpg")
    saltImage = salt(img, 5000)
    plt.imshow(saltImage)
    plt.show()