Cómo eliminar ruido de imágenes (rayos X) con Autoencoders

Hola! 🧡

Siguiendo con el tema de imágenes, hoy les voy a contar como eliminar el ruido de imágenes médicas, en este caso de mamografías tomadas de la base de datos mini-MIAS que se encuentra en el siguiente link http://peipa.essex.ac.uk/info/mias.html. En ella encontramos 322 imágenes de mamografías en resolución 1024x1024 pixeles.

El plan de trabajo será:

1) Cargar y explorar el dataset, también observar algunas de las imágenes en él

2) Dividir el conjunto de datos de imágenes en entrenamiento (80%) y prueba (20%)

3) Hacer aumento de los datos girando las fotografías

4) Agregar ruido Gaussiano a cada una de las imágenes en el entrenamiento y en prueba

5) Entrenar la red de Autoencoders Convolucional teniendo como entrada las imágenes con ruido y salida las imágenes sin ruido.

Este post está inspirado en una publicación titulada "Medical image denoising using convolutional denoising autoencoders" publicada en 2016 por Lovedeep Gondara Department of Computer Science, Simon Fraser University. En ella, también se usa el método de autoencoders para eliminar el ruido de imágenes médicas, obteniendo buenos resultados. He querido reproducir los resultados de estudio y mostrarles cómo lo he hecho en python usando las librerias Keras, sklearn y cv2.

Carga y Exploración de los Datos:

He hecho la carga de los datos en un cuaderno de Google Colab usando el siguiente código:

import tarfile
import shutil
import glob

# extracting all the files 
print('Extracting all the files now...') 
tar = tarfile.open('/content/all-mias.tar.gz', "r:gz")
tar.extractall()
tar.close()
print('Done!')

!cd /content

Ahora echemos un vistazo rápido de algunas de las imágenes en el dataset

# list all the image file names 
listdir=glob.glob('*.pgm')

import matplotlib.pyplot as plt
import cv2
%matplotlib inline

for i in range(9):
 # define subplot
    plt.subplot(330 + 1 + i)
 # define filename
    filename = '/content/' + str(listdir[i])
 # load image pixels
    img = cv2.imread(filename)
 # plot raw pixel data
    img=cv2.cvtColor(img, cv2.COLOR_BGR2GRAY ) 
    img=cv2.resize(img,(200,200))
    plt.imshow(img,cmap='gray')
    plt.axis('off')
# show the figure
plt.show()

Aumento de los datos

Vamos a rotar las imágenes haciendo uso de la función "flip" de la librería "cv2" para aumentar la cantidad de imágenes. Aumentar la cantidad de ejemplos que se presentan al entrenamiento generalmente mejora el desempeño de los modelos. He guardado las imágenes en una carpeta llamada "data".

import os
!mkdir data

for i in listdir:
    img=cv2.imread('/content/'+i)
    img=cv2.cvtColor(img, cv2.COLOR_BGR2GRAY ) 
    os.chdir('/content/data/')
    cv2.imwrite(i, img) 
 for j in range(-1,1):
        flip=cv2.flip(img, j)
        cv2.imwrite(str(j)+'_'+i, flip)
    os.chdir('/content/')
    os.remove(i) 

División de los datos en conjunto de entrenamiento y prueba

Ordenaremos nuestros datos en conjuntos de entrenamiento "train" y prueba "test", 80% para entrenamiento y 20% para pruebas.

Seleccionamos aleatoriamente las imágenes usando la librería sklearn de python, recordemos definir en "random state" por razones de reproducibilidad del trabajo.

os.chdir('/content/data/')
listdir=glob.glob('*.pgm')
os.chdir('/content/')

from sklearn.model_selection import train_test_split
trainlist,testlist=train_test_split(listdir,test_size=0.2,random_state=0)
print(len(trainlist),len(testlist))
772 194

Tenemos 772 imágenes para entrenamiento y 194 para pruebas seleccionadas aleatoriamente.

Adición del ruido

Definiremos una función que agregue ruido a las imágenes con la ayuda de la librería "numpy"

import numpy as np
def noisy(image):
    row,col= image.shape
    mean = 0
    sigma = 50
    gauss = np.random.normal(mean,sigma,(row,col))
    gauss = gauss.reshape(row,col)
    noisy = image + gauss
 return noisy

El resultado de añadir el ruido Gaussiano lo podemos ver en la imagen

img_noise=noisy( "gauss",img)
plt.imshow(img_noise,cmap='gray')

Nuestro objetivo es entrenar un modelo de Machine Learning con Redes Neurales Convolucionales que permita eliminar este ruido de las imagenes. Lo que haremos será añadir este ruido a todas las imágenes en nuestros datos, y entrenar el modelo mostrando como entrada del modelo la imagen con ruido y como salida la imagen sin ruido. De esta manera la red "aprenderá" a eliminar el ruido de las imágenes.

El siguiente código crea el set de datos como array de imágenes sin ruido para entrenamiento y prueba.

n=200
import numpy as np
train_nonoise = []
for filename in trainlist:
 # load image
    train = cv2.imread('data/' + filename)
    train = cv2.cvtColor(train, cv2.COLOR_BGR2GRAY ) 
    train=cv2.resize(train,(n,n))
 # store loaded image
    train_nonoise.append(train)
train_nonoise=np.array(train_nonoise)
print('loaded',  train_nonoise.shape[0], 'train no noised images' )


test_nonoise = []
for filename in testlist:
 # load image
    test = cv2.imread('data/' + filename)
    test = cv2.cvtColor(test, cv2.COLOR_BGR2GRAY )
    test=cv2.resize(test,(n,n)) 
 # store loaded image
    test_nonoise.append(test)
test_nonoise=np.array(test_nonoise)
print('loaded',  test_nonoise.shape[0], 'test no noised images' )

loaded 772 train no noised images loaded 194 test no noised images

Un código similar crea el array de imágenes con ruido para entrenamiento y prueba

train_gaussnoise = []
for filename in trainlist:
 # load image
    train = cv2.imread('data/' + filename)
    train = cv2.cvtColor(train, cv2.COLOR_BGR2GRAY ) 
    train=cv2.resize(train,(n,n))
    train=noisy( "gauss",train)
 # store loaded image
    train_gaussnoise.append(train)
train_gaussnoise=np.array(train_gaussnoise)
print('loaded',  train_gaussnoise.shape[0], 'train gauss noised images' )


test_gaussnoise = []
for filename in testlist:
 # load image
    test = cv2.imread('data/' + filename)
    test = cv2.cvtColor(test, cv2.COLOR_BGR2GRAY ) 
    test=cv2.resize(test,(n,n))
    test=noisy( "gauss",test)
 # store loaded image
    test_gaussnoise.append(test)
test_gaussnoise=np.array(test_gaussnoise)
print('loaded',  test_gaussnoise.shape[0], 'test noised gauss images' )

loaded 772 train gauss noised images 
loaded 194 test noised gauss images

Entrenamiento del modelo

Antes de entrenar el modelo vamos a convertir cada foto en un vector de una fila, creando una matriz de datos para entrenamiento de 200x200 columnas y 772 filas que corresponde al número de imágenes en el entrenamiento. Y para el conjunto de prueba tendremos una matriz de 200x200 columnas y 194 filas.

def preprocess(x):
    x = x.astype('float32') / 255.
 return x.reshape(-1, n, n, 1) 
X_train_noise = preprocess(train_gaussnoise)
X_train  = preprocess(train_nonoise)
X_test_noise = preprocess(test_gaussnoise)
X_test  = preprocess(test_nonoise)

Hemos normalizado las entradas de estas matrices dividiendo cada entrada por 255, recordemos que los colores en las imágenes están frecuentemente representados por números entre 0 y 255. Ahora si, vamos a crear el modelo de autoencoder para el entrenamiento:

from tensorflow.keras.layers import Input, Dense, Conv2D, MaxPool2D,MaxPooling2D, UpSampling2D
from tensorflow.keras.models import Model,Sequential
import tensorflow

def make_convolutional_autoencoder():
 # encoding
    input_layer = Input(shape=(n, n, 1))
 
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(input_layer)
    x = MaxPool2D( (2, 2), padding='same')(x)
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(x)
    x = MaxPool2D( (2, 2), padding='same')(x)
    x = Conv2D(16, (3, 3), activation='relu', padding='same')(x)
    latent_view    = MaxPool2D( (2, 2), padding='same')(x)

 # decoding architecture
    x = Conv2D(16, (3, 3), activation='relu', padding='same')(latent_view)
    x = UpSampling2D((2, 2))(x)
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(x)
    x = UpSampling2D((2, 2))(x)
    x = Conv2D(64, (3, 3), activation='relu',padding='same')(x)
    x = UpSampling2D((2, 2))(x)
    output_layer   = Conv2D(1, (3, 3), padding='same', activation='sigmoid')(x)

 # autoencoder
    autoencoder = Model(input_layer, output_layer)
    autoencoder.compile(optimizer='adam', 
                        loss='binary_crossentropy')
 return autoencoder
# create a convolutional autoencoder
autoencoder = make_convolutional_autoencoder()
autoencoder.summary()

Con esta arquitectura de autoencoder hemos entrenado nuestro conjunto de datos

autoencoder.fit(X_train_noise, X_train, 
                epochs=100, 
                batch_size=60, 
                validation_data=(X_test_noise, X_test))

Nuestro último paso será contrastar nuestros resultados con la imagen original y la imagen con ruido.

X_test_pred = autoencoder.predict(X_test_noise)
plt.imshow(cv2.flip(X_test_pred.reshape(X_test.shape[0],n,n)[5],1),cmap='gray')

Como podemos observar, hemos eliminado el ruido de la imagen, sin embargo no hemos podido recuperar la nitidez de la imagen original. Podemos ver que fue posible recuperar algunos patrones que se observan en la imagen original así como también la forma de la mama. Fué posible reproducir los resultados del paper haciendo uso de la misma arquitectura de autoencoders.

Nos vemos!