Шумоподавление автоэнкодеров

В одном из моих предыдущих постов я рассказывал о разных типах автоэнкодеров. В этом посте я уделю больше внимания одному конкретному типу — шумоподавляющим автоэнкодерам.

Удаление шума с изображений является сложной и важной задачей в области обработки изображений и компьютерного зрения. Удаление шума может быть весьма полезным при анализе изображений. Целью, вероятно, является извлечение наиболее важных функций, и если у вас есть способ выбросить всю ненужную информацию, желательно.

Давайте посмотрим на пример, где автоэнкодер смог убрать шум. Из этого:

Сеть была обучена делать это:

Идея состоит в том, чтобы исказить входное изображение, но установить исходное изображение в качестве желаемого результата. Таким образом, наша сеть не просто выучит функцию идентификации, потому что вход и выход разные.

Давайте углубимся в реализацию шумоподавляющего автоэнкодера. Для набора данных я выбрал файл mnist.

Сначала импортируем необходимые библиотеки:

import numpy  as np
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import Dense
import PIL
import tensorflow as tf
from keras.datasets import mnist
from tensorflow import keras
from keras.models import Sequential
from keras.layers import (
Conv2DTranspose,Reshape,BatchNormalization, SeparableConv2D, MaxPooling2D, Activation, Flatten, Dropout, Dense, Conv2D,UpSampling2D,Input
)
from PIL import Image as im
from keras.models import Model
from sklearn.model_selection import train_test_split
import cv2

Затем определите функцию для загрузки данных и нормализуйте изображения:

def load_data():
     (x_train, y_train), (x_test, y_test) = mnist.load_data()
     m = x_train.shape[0]
     n = x_test.shape[0]
     x_train = np.reshape(x_train, (-1,28,28,1))
     x_test = np.reshape(x_test,(-1,28,28,1))
     y_train = tf.keras.utils.to_categorical(y_train) 
     y_test = tf.keras.utils.to_categorical(y_test)
     x_train = x_train.astype('float32')
     x_test = x_test.astype('float32')
     x_train /= 255
     x_test /= 255
     return x_train,y_train,x_test,y_test

И загрузите данные:

x_train,y_train,x_test,y_test = load_data()

Мы хотим добавить к входным данным шум, поэтому определяем функцию Gaussian_noise:

def add_gaussian_noise(X_imgs):
  gaussian_noise_imgs = []
  row, col, ch= X_imgs[0].shape
  for X_img in X_imgs:
     gaussian = X_img + 0.3*np.random.normal(loc = 0.0,scale =    1.0,size = (row, col, 1))
     gaussian_img = np.clip(gaussian, 0., 1.)
     gaussian_noise_imgs.append(gaussian_img)
     gaussian_noise_imgs = np.array(gaussian_noise_imgs, dtype =   np.float32)
  return gaussian_noise_imgs

Добавьте шум к входным изображениям:

x_train_noisy = add_gaussian_noise(x_train)
x_test_noisy = add_gaussian_noise(x_test)

Давайте посмотрим, как наши искаженные данные выглядят по сравнению с оригиналом:

Создайте кодировщик для нашего автоэнкодера:

def encoder(input_img):
  conv1_1 = Conv2D(16, (3, 3), activation='relu', padding='same')(input_img)
  conv1_2 = BatchNormalization()(conv1_1)
  conv1_3 = MaxPooling2D((2,2))(conv1_2)
  conv2_1 =   Conv2D(32,kernel_size=3,strides=2,padding='same',activation='relu')(conv1_3)
  conv2_2 = BatchNormalization()(conv2_1)
  conv2_3 = MaxPooling2D((2,2))(conv2_2)
  conv3_1 =  Conv2D(64,kernel_size=3,strides=2,padding='same',activation='relu')(conv2_3)
  conv3_2 = BatchNormalization()(conv3_1)
  conv4 = Flatten()(conv3_2)
  conv4 = Dense(576)(conv4)
  return conv4

И создайте декодер:

def decoder(encode):
  conv5_1 = Reshape((3,3,64))(encode)
  conv5_2 = Conv2DTranspose(32, kernel_size = 3, strides = 2,activation='relu', padding='valid')(conv5_1)
  conv6_1 = Conv2DTranspose(16, kernel_size = 3, strides = 2,activation='relu', padding='same')(conv5_2)
  conv7_1 = Conv2DTranspose(1, kernel_size=3, strides=2, padding='same', activation="sigmoid")(conv6_1)
  return conv7_1

Определите нашу модель:

inChannel = 1
x, y = 28, 28
input_img = Input(shape = (x, y,inChannel))
autoencoder = Model(input_img, decoder(encoder(input_img)))

Скомпилируйте нашу модель. Здесь я использовал оптимизатор SGD:

opt = tf.keras.optimizers.SGD(learning_rate=0.5)
autoencoder.compile(loss='binary_crossentropy', optimizer = opt)

Разделите тренировочные данные:

train_X,valid_X,train_ground,valid_ground =
train_test_split(x_train_noisy,x_train,test_size=0.2,random_state=13)

Обучите модель и сохраните веса:

autoencoder_train = autoencoder.fit(train_X, train_ground, batch_size=batch_size,epochs=epochs,verbose=1,validation_data=(valid_X, valid_ground))
autoencoder.save_weights('autoencoder.h5')

После тренировки посмотрите на результаты:

reconstructs = autoencoder.predict(x_test_noisy)
reconstructs = np.reshape(reconstructs,(-1,28,28))
x_test_noisy = np.reshape(x_test_noisy,(-1,28,28))
plt.figure(figsize=(15,12))
for i in range(5):
  plt.subplot(2,5,1+i)
  plt.imshow(x_test_noisy[i])
  plt.axis('off')
  plt.subplot(2,5,6+i)
  plt.imshow(reconstructs[i])
  plt.axis('off')
  plt.show()

Шум успешно удален. Результаты потрясающие, и я действительно советую поиграть с разными архитектурами и наборами данных.

Также попробуйте использовать полученные результаты для какой-нибудь классификации, например. Однако имейте в виду, что шумоподавляющие автоэнкодеры хорошо работают с данными, которые незначительно отличаются от обученных данных. Например, если бы я попробовал эту модель в наборе данных ImageNet, я мог бы также получить большую ошибку реконструкции.

Надеюсь, этот пост был полезен :)