import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

img = cv.imread("../SampleImages/pomeranian.png", cv.IMREAD_COLOR)
print(img.shape)
pixel_count = img.shape[0] * img.shape[1]
print(pixel_count)


#为图像添加椒盐噪声
#参考资料：https://blog.csdn.net/sinat_29957455/article/details/123977298
img_pepper_salt_noise = img.copy()
#椒盐噪声中盐（白点）和椒（黑点）的比例
salt_ratio = 0.5
pepper_ratio = 1 - salt_ratio
#噪点占图像像素比例
noise_ratio = 0.03
#添加salt噪声
num_salt = np.ceil(noise_ratio * pixel_count * salt_ratio)
#噪声添加位置
noise_locations = [np.random.randint(0, i - 1, int(num_salt)) for i in img.shape]
img_pepper_salt_noise[noise_locations[0],noise_locations[1],:] = [255,255,255]
#添加pepper噪声
num_pepper = np.ceil(noise_ratio * pixel_count * pepper_ratio)
noise_locations = [np.random.randint(0, i - 1, int(num_salt)) for i in img.shape]
img_pepper_salt_noise[noise_locations[0],noise_locations[1],:] = [0,0,0]

#图像中值滤波
#cv.medianBlur(src, ksize)
#src: 待处理图像
#ksize：kernel大小，核大小必须是比1大的奇数，比如3、5、7等
#参考资料：https://blog.csdn.net/qq_49478668/article/details/123485382
img_blur = cv.medianBlur(img_pepper_salt_noise, 3)

#显示图像
fig,axes = plt.subplots(nrows=1, ncols=3, figsize=(15,15), dpi=100)
axes[0].imshow(img[:,:,::-1])
axes[0].set_title("Original");
axes[1].imshow(img_pepper_salt_noise[:,:,::-1])
axes[1].set_title("Salt And Pepper Noise");
axes[2].imshow(img_blur[:,:,::-1])
axes[2].set_title("Median Blurred");