使用pytorch搭建自编码器，实现图像的去噪

import torchvision import matplotlib.pyplot as plt from torch.utils.data import DataLoader import numpy as np import random import PIL.Image as Image import torchvision.transforms as transforms class AddPepperNoise(object): """增加椒盐噪声 Args: snr （float）: Signal Noise Rate p (float): 概率值，依概率执行该操作 """ def __init__(self, snr, p=0.9): assert isinstance(snr, float) and (isinstance(p, float)) self.snr = snr self.p = p def __call__(self, img): """ Args: img (PIL Image): PIL Image Returns: PIL Image: PIL image. """ if random.uniform(0, 1) < self.p: img_ = np.array(img).copy() h, w = img_.shape signal_pct = self.snr noise_pct = (1 - self.snr) mask = np.random.choice((0, 1, 2), size=(h, w), p=[signal_pct, noise_pct/2., noise_pct/2.]) img_[mask == 1] = 255 # 盐噪声 img_[mask == 2] = 0 # 椒噪声 return Image.fromarray(img_.astype('uint8')) else: return img class Gaussian_noise(object): """增加高斯噪声 此函数用将产生的高斯噪声加到图片上 传入: img : 原图 mean : 均值 sigma : 标准差 返回: gaussian_out : 噪声处理后的图片 """ def __init__(self, mean, sigma): self.mean = mean self.sigma = sigma def __call__(self, img): """ Args: img (PIL Image): PIL Image Returns: PIL Image: PIL image. """ # 将图片灰度标准化 img_ = np.array(img).copy() img_ = img_ / 255.0 # 产生高斯 noise noise = np.random.normal(self.mean, self.sigma, img_.shape) # 将噪声和图片叠加 gaussian_out = img_ + noise # 将超过 1 的置 1，低于 0 的置 0 gaussian_out = np.clip(gaussian_out, 0, 1) # 将图片灰度范围的恢复为 0-255 gaussian_out = np.uint8(gaussian_out*255) # 将噪声范围搞为 0-255 # noise = np.uint8(noise*255) return Image.fromarray(gaussian_out) if __name__=='__main__': train_datasets = torchvision.datasets.MNIST('./', train=True, download=True) test_datasets = torchvision.datasets.MNIST('./', train=False, download=True) print('训练集的数量', len(train_datasets)) print('测试集的数量', len(test_datasets)) train_loader = DataLoader(train_datasets, batch_size=100, shuffle=True) test_loader = DataLoader(test_datasets, batch_size=1, shuffle=False) transform = transforms.Compose([ transforms.ToPILImage(), Gaussian_noise(0, 0.1), AddPepperNoise(0.9) # transforms.ToTensor() ]) print('训练集可视化') fig = plt.figure() for i in range(12): plt.subplot(3, 4, i + 1) img = train_datasets.train_data[i] label = train_datasets.train_labels[i] # noise = np.random.normal(0.1, 0.1, img.shape) # img=transform(img) plt.imshow(img, cmap='gray') plt.title(label) plt.xticks([]) plt.yticks([]) plt.show()