¿Es posible diagnosticar el COVID-19 con Deep Learning?: «Tosa, por favor»

Identificar sonidos es una tarea que el ser humano hace de manera muy eficiente y, durante los últimos años, los ordenadores están ganando cada vez más experiencia en la misma tarea.

Con la emergencia de la COVID-19 están surgiendo una varios proyectos que intentan identificar si una persona está afectada por esta enfermedad; uno de los ejemplos es hacer uso de las toses y/o sibilancias respiratorias tanto de sujetos afectados y sanos para poder hacer una clasificación. Conseguir un gran número de muestras es algo imprescindible para poder encaminar la consecución de un modelo robusto.

¿Cómo se podría abordar un problema de este tipo? Una de las maneras de discretizar la naturaleza u origen de los sonidos es mediante clasificación de imágenes. Para ello, en esta prueba de concepto vamos a hacer uso de la potencia de  modelos convolucionales ya entrenados y, mediante fine-tuning, conseguir buenas precisiones en nuestra tarea ahorrando tiempo de entrenamiento. Esto nos permite no sólo abordar sonidos relacionados con los efectos del SARS-CoV-2 en el sistema respiratorio, sino cualquier sonido convenientemente etiquetado.

En busca de un dataset

Nuestra meta, como hemos comentado, es clasificar sonidos en base a unas etiquetas ya asignadas. El problema actual reside en la ausencia de datasets. No contamos con un dataset de calidad, única y exclusivamente compuesto por toses, que nos permita determinar la afección que sufre una persona.

El único dataset que tenemos disponible hasta ahora es uno liberado por la farmacéutica Pfizer, llamado Sick Sounds, que se puede encontrar en el siguiente enlace.

Este dataset separa a los audios entre personas enfermas y no enfermas en dos directorios distintos. Sin embargo cuando uno profundiza en los audios obtenidos, se ve que las clases están un poco desbalanceadas y que los audios catalogados como «sujeto no enfermo» en realidad no son audios de toses, sino de música, risas, conversaciones y un largo etcétera de diferentes situaciones.

A pesar de esta inconveniencia, podemos tomar este dataset como una primera referencia para intentar ver si conseguimos una buena precisión clasificando audios usando Redes Neuronales Convolucionales.

Extracción de características

Hay muchas características qué podemos usar para entrenar un modelo. En lo referente al campo de speech recognition, podemos usar los Mel-Frequency Cepstral Coefficients (MFCC). Lo mejor de este tipo de coeficientes es que nos dan una buena representación del audio original.

Sería conveniente también tener más información más allá de las características obtenidas de MFCC, pero las redes llamadas WaveNet pueden ser muy difíciles de entrenar y también de ejecutar. Por ello vamos a usar espectrogramas.

Espectrogramas

En investigación de audios, un espectrograma se considera como una representación gráfica de un audio que tiene la frecuencia en el eje de ordenadas y el tiempo en eje de coordenadas y una tercera dimensión – con colores – que representa la intensidad del sonido para cada tiempo y cada frecuencia. Aquí tenemos un ejemplo de un espectrograma.

Lo mejor del uso de los espectrogramas es que nos permite cambiar de un problema de reconocimiento de audio a un problema de reconocimiento de imágenes. Pasar de audio a imagen es bastante sencillo. Aquí tenemos un simple script que nos permite obtener el espectrograma de un audio, usando la librería librosa.

Fine-tuning con Redes Neuronales Convolucionales

Una vez que hemos conseguido que nuestros audios queden representados como imágenes, podemos proceder a su clasificación haciendo uso de redes neuronales. Para ello vamos a usar una red convolucional ya entrenada.

Si tuviéramos que entrenar una red convolucional desde cero nos llevaría bastante tiempo realizar está tarea y dada la calidad de nuestro dataset con total seguridad acabaríamos con un problema de sobreajuste en nuestros datos.

Para ello vamos a usar Transfer Learning, el cual nos permitirá hacer uso de una red neuronal ya entrenada con datasets similares y re-entrenar unas pocas capas densas al final de toda la parte convolucional, basadas en nuevas categorías.

Google lanzó en su día un modelo presentado llamado Inception entrenado para clasificar imágenes del ImageNet dataset. De hecho la gente de Tensorflow facilita un script para poder hacer un fine-tuning de inserción con nuevas categorías.

Para obtener este script simplemente ejecutamos el siguiente comando:

$ curl -O \
https://raw.githubusercontent.com/tensorflow/tensorflow/r1.1/tensorflow/examples/image_retraining/retrain.py

Una vez que tengamos todas nuestras imágenes de espectrogramas ordenadas en carpetas en función de si el sujeto está enfermo o no, podemos lanzar nuestro re-entrenamiento:

$ python retrain.py \
  --bottleneck_dir=bottlenecks \
  --how_many_training_steps=40000 \
  --model_dir=inception \
  --summaries_dir=training_summaries/basic \
  --output_graph=retrained_graph.pb \
  --output_labels=retrained_labels.txt \
  --image_dir=melspectrograms/train

Tras haber entrenado nuestro modelo durante una media hora, se ha obtenido una precisión aproximada del 78%. La precisión, como puede verse, es mejorable.

Usemos nuestro modelo: construyendo un frontal web

Se ha trabajado en la construcción de un Frontal Web muy rudimentario que nos permita servir a través de un navegador audios que nos permita dar los siguientes pasos:

1. Transformar audios a imágenes (espectrogramas);
2. Testearlos contra nuestra red neuronal;
3. Obtener por terminal un score de salud del individuo con la estructura.

{
    "Diagnóstico": "sick" | "not sick", 
    "score": any float in interval 0.0 - 1.0
}

La estructura de la prueba de frontal es la siguiente:

|-- melspectrograms
|   |-- train
|       |-- sick
|           |-- ... .png
|       |-- not sick 
|       	|-- ... .png
|   |-- validation
|       |-- sick
|	        |-- ... .png
|	    |-- not sick 
|	        |-- ... .png
|-- inception
|   |-- classify_image_graph_def.pb
|-- bottlenecks
|-- static
|-- templates
|-- app.py
|-- retrain.py
|-- retrained_graph.pb
`-- retrained_labels.txt

Nos vamos a centrar en el script app.py nos permitirá entender cómo procesar un audio, transformarlo a imagen y servirlo a nuestro modelo en tiempo real.

from flask import Flask
from flask import request
from flask import render_template, jsonify
import os, sys
import librosa
import numpy as np
import matplotlib.pyplot as plt
import librosa.display
import tensorflow as tf
import logging
logging.getLogger('tensorflow').disabled = True

app = Flask(__name__)

Definimos la única ruta de nuestra simple webapp:

@app.route("/", methods=['POST', 'GET'])
def index():
    # Unpersists graph from file
    # Cargamos el modelo
    with tf.gfile.FastGFile("retrained_graph.pb", 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        tf.import_graph_def(graph_def, name='')
    if request.method == "POST":
        with tf.Session() as sess:
            # Alimentamos image_data como input al grafo para predicción
            softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
            f = open('./file.wav', 'wb')
            f.write(request.get_data("audio_data"))
            f.close()
            # Procesamos el dato de audio facilitado al frontal
            y, sr = librosa.load(request.files['audio_data'])

            # Let's make and display a mel-scaled power (energy-squared) spectrogram
            S = librosa.feature.melspectrogram(y, sr=sr, n_mels=128)

            # Convertimos a log scale (dB). utilizaremos la potencia de pico como referencia.
            log_S = librosa.amplitude_to_db(S, ref=np.max)

Ahora podemos guardar, para cada muestra de audio que se suba al frontal, una imagen de frecuencias.

            # Make a new figure
            fig = plt.figure(figsize=(12,4))
            ax = plt.Axes(fig, [0., 0., 1., 1.])
            ax.set_axis_off()
            fig.add_axes(ax)

            # Display the spectrogram on a mel scale
            librosa.display.specshow(log_S, sr=sr, x_axis='time', y_axis='mel')

            # Make the figure layout compact
            image_path = 'testing/tmp.png'
            plt.savefig(image_path)
            plt.close()

Ahora, con la imagen generada, sólo tenemos que testearla contra el modelo y obtener una etiqueta en base a lo aprendido durante el entrenamiento.

            # Read in the image_data
            image_data = tf.gfile.FastGFile(image_path, 'rb').read()

            # Loads label file, strips off carriage return
            label_lines = [line.rstrip() for line
                            in tf.gfile.GFile("retrained_labels.txt")]

            predictions = sess.run(softmax_tensor, \
                    {'DecodeJpeg/contents:0': image_data})

            # Sort to show labels of first prediction in order of confidence
            top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]
            
            payload = {
                "diagnosis": label_lines[top_k[0]], \
                "score": predictions[0][top_k[0]]
            }
            # Mostramos por consola el resultado del diagnóstico
            print(payload)
    
        return render_template('index.html', request="POST")
    else:
        return render_template('index.html')

if __name__ == "__main__":
    app.run(debug=True, \
            threaded=True, \
            port="5000")

Para iniciar la webapp, dentro del entorno virtual correcto, símplemente ejecutamos:

$ python app.py

Tras ejecutar este comando se levantará un servidor al que se puede acceder a través de la ruta http://localhost:5000/

Aquí se incluye un vídeo con una breve demo.

Conclusiones

Con unos conocimientos básicos de Python, Tensorflow y de Transfer Learning hemos podido clasificar un dataset de sonidos de manera rápida con una precisión cercana al 80%.

Con un modelo ya entrenado, hemos esbozado un frontal muy simple que nos permite hacer servir grabaciones desde el micrófono de nuestro navegador a nuestro motor para poder predecir una score de enfermedad en base a lo obtenido tras el entrenamiento.

El código se encuentra disponible en el siguiente repositorio de GitHub de LUCA.

Siguientes pasos

• Se tendría que encontrar un dataset real de toses y sonidos respiratorios que nos permita construir un modelo de mejor calidad.
• Trabajar en un mejor fine-tuning del modelo para mejorar la precisión.
• Investigar WaveNet como una opción para entrenar un nuevo modelo.

Para mantenerte al día con LUCA visita nuestra página web,  suscríbete a LUCA Data Speaks o síguenos en Twitter, LinkedIn y YouTube.