La Guía Definitiva de nnU-Net
La Guía Definitiva de nnU-Net' (The Ultimate Guide to nnU-Net)
Todo lo que necesitas saber para entender el Estado del Arte nnU-Net y cómo aplicarlo a tu propio conjunto de datos.
Durante mi pasantía de investigación en Aprendizaje Profundo y Neurociencias en la Universidad de Cambridge, utilicé mucho nnU-Net, que es una referencia extremadamente sólida en Segmentación Semántica de Imágenes.
Sin embargo, tuve dificultades para comprender completamente el modelo y cómo entrenarlo, y no encontré mucha ayuda en Internet. Ahora que estoy cómodo con esto, he creado este tutorial para ayudarte, ya sea en tu búsqueda por entender mejor lo que hay detrás de este modelo o cómo usarlo en tu propio conjunto de datos.
A lo largo de esta guía, aprenderás:
- Desarrollar una visión concisa de las principales contribuciones de nnU-Net.
- Aprender cómo aplicar nnU-Net a tu propio conjunto de datos.
Todo el código está disponible en este cuaderno de Google Collab
Este trabajo me llevó una cantidad significativa de tiempo y esfuerzo. Si encuentras este contenido valioso, considera seguirme para aumentar su visibilidad y ayudar a apoyar la creación de más tutoriales como este.
Una Breve Historia de nnU-Net
Reconocido como un modelo de vanguardia en Segmentación de Imágenes, nnU-Net es una fuerza indomable cuando se trata de procesamiento de imágenes 2D y 3D. Su rendimiento es tan sólido que sirve como una referencia sólida con la que se comparan nuevas arquitecturas de visión por computadora. En esencia, si te estás aventurando en el mundo del desarrollo de modelos de visión por computadora novedosos, considera a nnU-Net como tu ‘objetivo a superar’.
Esta poderosa herramienta se basa en el modelo U-Net (Puedes encontrar uno de mis tutoriales aquí: Preparar tu primer U-Net), que hizo su debut en 2015. La denominación “nnU-Net” significa “No New U-Net”, en referencia al hecho de que su diseño no introduce alteraciones arquitectónicas revolucionarias. En cambio, toma la estructura existente de U-Net y exprime todo su potencial utilizando un conjunto de estrategias de optimización ingeniosas.
A diferencia de muchas redes neuronales modernas, nnU-Net no se basa en conexiones residuales, conexiones densas o mecanismos de atención. Su fortaleza radica en su estrategia de optimización meticulosa, que incluye técnicas como el remuestreo, la normalización, la elección juiciosa de la función de pérdida, la configuración del optimizador, la ampliación de datos, la inferencia basada en parches y la combinación de modelos. Este enfoque integral permite a nnU-Net empujar los límites de lo que se puede lograr con la arquitectura original de U-Net.
Explorando Arquitecturas Diversas dentro de nnU-Net
Aunque pueda parecer una entidad singular, nnU-Net es en realidad un término genérico que representa tres tipos distintos de U-Nets:
- U-Net 2D: Posiblemente la variante más conocida, opera directamente en imágenes 2D.
- U-Net 3D: Esta es una extensión de U-Net 2D y es capaz de manejar imágenes 3D directamente mediante la aplicación de convoluciones 3D.
- U-Net en Cascada: Este modelo genera segmentaciones de baja resolución y posteriormente las mejora.
Cada una de estas arquitecturas aporta sus fortalezas únicas y, inevitablemente, tiene ciertas limitaciones.
Por ejemplo, emplear un U-Net 2D para la segmentación de imágenes en 3D puede parecer poco intuitivo, pero en la práctica puede ser altamente efectivo. Esto se logra al dividir el volumen 3D en planos 2D.
Aunque un U-Net 3D puede parecer más sofisticado debido a su mayor número de parámetros, no siempre es la solución más eficiente. En particular, los U-Net 3D suelen tener dificultades con la anisotropía, que ocurre cuando las resoluciones espaciales difieren a lo largo de diferentes ejes (por ejemplo, 1 mm a lo largo del eje x y 1.2 mm a lo largo del eje z).
La variante U-Net Cascade resulta especialmente útil al lidiar con imágenes grandes. Emplea un modelo preliminar para condensar la imagen, seguido de un U-Net 3D estándar que produce segmentaciones de baja resolución. Las predicciones generadas se escalan posteriormente, lo que resulta en una salida refinada y completa.
Típicamente, la metodología implica entrenar las tres variantes del modelo dentro del marco de trabajo de nnU-Net. El siguiente paso puede ser elegir el mejor rendimiento entre las tres o emplear técnicas de ensamblaje. Una técnica de ensamblaje podría involucrar la integración de las predicciones tanto de los U-Net 2D como de los U-Net 3D.
Sin embargo, es importante tener en cuenta que este procedimiento puede ser bastante lento (y también costoso debido a que se necesitan créditos de GPU). Si tus restricciones solo permiten el entrenamiento de un solo modelo, no te preocupes. Puedes elegir entrenar solo un modelo, ya que el modelo de ensamblaje solo aporta mejoras muy marginales.
Esta tabla ilustra la variante del modelo con mejor rendimiento en relación a conjuntos de datos específicos:
Adaptación dinámica de topologías de red
Dadas las discrepancias significativas en el tamaño de las imágenes (considera la forma mediana de 482 × 512 × 512 para las imágenes del hígado frente a 36 × 50 × 35 para las imágenes del hipocampo), nnU-Net se adapta inteligentemente al tamaño del parche de entrada y al número de operaciones de pooling por eje. Esto implica ajustar automáticamente la cantidad de capas convolucionales por conjunto de datos, facilitando la agregación efectiva de información espacial. Además de adaptarse a las geometrías de imagen variadas, este modelo tiene en cuenta restricciones técnicas, como la memoria disponible.
Es crucial tener en cuenta que el modelo no realiza la segmentación directamente en toda la imagen, sino en parches cuidadosamente extraídos con regiones superpuestas. Las predicciones en estos parches se promedian posteriormente, lo que lleva a la salida final de segmentación.
Pero tener un parche grande significa un mayor uso de memoria, y el tamaño del lote también consume memoria. El compromiso adoptado es siempre priorizar el tamaño del parche (la capacidad del modelo) en lugar del tamaño del lote (solo útil para la optimización).
Aquí está el algoritmo heurístico utilizado para calcular el tamaño óptimo del parche y el tamaño del lote:
Y esto es cómo se ve para diferentes conjuntos de datos y dimensiones de entrada:
¡Genial! Ahora repasemos rápidamente todas las técnicas utilizadas en nnU-Net:
Entrenamiento
Todos los modelos se entrenan desde cero y se evalúan utilizando validación cruzada de cinco pliegues en el conjunto de entrenamiento, lo que significa que el conjunto de datos de entrenamiento original se divide aleatoriamente en cinco partes iguales, o ‘pliegues’. En este proceso de validación cruzada, cuatro de estos pliegues se utilizan para el entrenamiento del modelo y el pliegue restante se utiliza para la evaluación o prueba. Este proceso se repite cinco veces, utilizando cada uno de los cinco pliegues exactamente una vez como conjunto de evaluación.
Para la pérdida, utilizamos una combinación de la Pérdida de Dado y la Entropía Cruzada. Esta es una pérdida muy frecuente en la Segmentación de Imágenes. Más detalles sobre la Pérdida de Dado en V-Net, el hermano mayor de U-Net
Técnicas de aumento de datos
El nnU-Net tiene un sólido pipeline de aumento de datos. Los autores utilizan rotaciones aleatorias, escalado aleatorio, deformación elástica aleatoria, corrección gamma y espejado.
NB: Puedes añadir tus propias transformaciones modificando el código fuente
Inferencia basada en parches
Así como dijimos, el modelo no predice directamente sobre la imagen de resolución completa, lo hace en los parches extraídos y luego agrega la predicción.
Esto es cómo se ve:
NB: Los parches en el centro de la imagen tienen más peso que los que están en los lados, porque contienen más información y el modelo funciona mejor en ellos
Ensamblaje de modelos en pares
Entonces, si recuerdas bien, podemos entrenar hasta 3 modelos diferentes, 2D, 3D y en cascada. Pero cuando hacemos inferencia solo podemos usar un modelo a la vez ¿verdad?
Bueno, resulta que no, diferentes modelos tienen diferentes fortalezas y debilidades. Así que en realidad podemos combinar las predicciones de varios modelos para que si un modelo está muy seguro, prioricemos su predicción.
nnU-Net prueba todas las combinaciones de 2 modelos entre los 3 modelos disponibles y selecciona la mejor.
En la práctica, hay 2 formas de hacerlo:
Votación dura: Para cada píxel, miramos todas las probabilidades generadas por los 2 modelos y tomamos la clase con la probabilidad más alta.
Votación suave: Para cada píxel, promediamos la probabilidad de los modelos y luego tomamos la clase con la máxima probabilidad.
Implementación práctica
Antes de empezar, puedes descargar el conjunto de datos aquí y seguir el cuaderno de Google Collab.
Si no entendiste nada de la primera parte, no te preocupes, esta es la parte práctica, solo necesitas seguirme y aún así obtendrás los mejores resultados.
Necesitas una GPU para entrenar el modelo, de lo contrario no funciona. Puedes hacerlo localmente o en Google Collab, no olvides cambiar el entorno de ejecución > GPU
Entonces, en primer lugar, debes tener un conjunto de datos listo con imágenes de entrada y su correspondiente segmentación. Puedes seguir mi tutorial descargando este conjunto de datos listo para la segmentación cerebral 3D, y luego puedes reemplazarlo con tu propio conjunto de datos.
Descargando datos
En primer lugar, debes descargar tus datos y colocarlos en la carpeta de datos, nombrando las dos carpetas “input” y “ground_truth” que contienen la segmentación.
Para el resto del tutorial, utilizaré el conjunto de datos MindBoggle para la segmentación de imágenes. Puedes descargarlo en este Google Drive:
Se nos dan exploraciones de resonancia magnética 3D del cerebro y queremos segmentar la materia blanca y gris:
Debería verse así:
Configuración del directorio principal
Si ejecutas esto en Google Colab, establece collab = True, de lo contrario, collab = False
collab = Trueimport osimport shutil#libreríasfrom collections import OrderedDictimport jsonimport numpy as np#visualización del conjunto de datosimport matplotlib.pyplot as pltimport nibabel as nibif collab: from google.colab import drive drive.flush_and_unmount() drive.mount('/content/drive', force_remount=True) # Cambia "neurosciences-segmentation" por el nombre de tu carpeta de proyecto root_dir = "/content/drive/MyDrive/neurosciences-segmentation"else: # obtén el directorio del directorio padre root_dir = os.getcwd()input_dir = os.path.join(root_dir, 'data/input')segmentation_dir = os.path.join(root_dir, 'data/ground_truth')my_nnunet_dir = os.path.join(root_dir,'my_nnunet')print(my_nnunet_dir)Ahora vamos a definir una función que crea carpetas para nosotros:
def make_if_dont_exist(folder_path,overwrite=False): """ crea una carpeta si no existe entrada: folder_path : ruta relativa de la carpeta que necesita ser creada over_write :(por defecto: False) si es True, sobrescribe la carpeta existente """ if os.path.exists(folder_path): if not overwrite: print(f'{folder_path} existe.') else: print(f"{folder_path} sobrescrita") shutil.rmtree(folder_path) os.makedirs(folder_path) else: os.makedirs(folder_path) print(f"{folder_path} creada!")Y usamos esta función para crear nuestra carpeta “my_nnunet” donde se guardará todo
os.chdir(root_dir)make_if_dont_exist('my_nnunet', overwrite=False)os.chdir('my_nnunet')print(f"Directorio de trabajo actual: {os.getcwd()}")Instalación de bibliotecas
Ahora vamos a instalar todos los requisitos. Primero, vamos a instalar la biblioteca nnunet. Si estás en un cuaderno, ejecuta esto en una celda:
!pip install nnunetDe lo contrario, puedes instalar nnunet directamente desde la terminal con
pip install nnunetAhora vamos a clonar el repositorio de git de nnUnet y NVIDIA apex. Esto contiene los scripts de entrenamiento, así como un acelerador de GPU.
!git clone https://github.com/MIC-DKFZ/nnUNet.git!git clone https://github.com/NVIDIA/apex# repository dir is the path of the github folderrespository_dir = os.path.join(my_nnunet_dir,'nnUNet')os.chdir(respository_dir)!pip install -e!pip install --upgrade git+https://github.com/nanohanno/hiddenlayer.git@bugfix/get_trace_graph#egg=hiddenlayerCreación de las carpetas
nnUnet requiere una estructura de carpetas muy específica.
task_name = 'Tarea001' #cambia aquí para un nombre de tarea diferente# Definimos todas las rutas necesariasnnunet_dir = "nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data"task_folder_name = os.path.join(nnunet_dir,task_name) train_image_dir = os.path.join(task_folder_name,'imagesTr') # ruta a las imágenes de entrenamientotrain_label_dir = os.path.join(task_folder_name,'labelsTr') # ruta a las etiquetas de entrenamientotest_dir = os.path.join(task_folder_name,'imagesTs') # ruta a las imágenes de pruebamain_dir = os.path.join(my_nnunet_dir,'nnUNet/nnunet') # ruta al directorio principaltrained_model_dir = os.path.join(main_dir, 'nnUNet_trained_models') # ruta a los modelos entrenadosOriginalmente, el nnU-Net fue diseñado para un desafío de decatlón con diferentes tareas. Si tienes diferentes tareas, simplemente ejecuta esta celda para todas tus tareas.
# Creación de todas las carpetasoverwrite = False # Establece esto como True si quieres sobrescribir las carpetasmake_if_dont_exist(task_folder_name,overwrite = overwrite)make_if_dont_exist(train_image_dir, overwrite = overwrite)make_if_dont_exist(train_label_dir, overwrite = overwrite)make_if_dont_exist(test_dir,overwrite= overwrite)make_if_dont_exist(trained_model_dir, overwrite=overwrite)Ahora deberías tener una estructura como esta:
Configuración de las variables de entorno
El script necesita saber dónde colocaste tus datos sin procesar, dónde puede encontrar los datos preprocesados, y dónde debe guardar los resultados.
os.environ['nnUNet_raw_data_base'] = os.path.join(main_dir,'nnUNet_raw_data_base')os.environ['nnUNet_preprocessed'] = os.path.join(main_dir,'preprocessed')os.environ['RESULTS_FOLDER'] = trained_model_dirMover los archivos en los repositorios correctos:
Definimos una función que moverá nuestras imágenes a los repositorios correctos en la carpeta nnunet:
def copy_and_rename(old_location,old_file_name,new_location,new_filename,delete_original = False): shutil.copy(os.path.join(old_location,old_file_name),new_location) os.rename(os.path.join(new_location,old_file_name),os.path.join(new_location,new_filename)) if delete_original: os.remove(os.path.join(old_location,old_file_name))Ahora ejecutemos esta función para las imágenes de entrada y de referencia:
list_of_all_files = os.listdir(segmentation_dir)list_of_all_files = [file_name for file_name in list_of_all_files if file_name.endswith('.nii.gz')]for file_name in list_of_all_files: copy_and_rename(input_dir,file_name,train_image_dir,file_name) copy_and_rename(segmentation_dir,file_name,train_label_dir,file_name)Ahora tenemos que renombrar los archivos para que sean aceptados por el formato nnUnet, por ejemplo subject.nii.gz se convertirá en subject_0000.nii.gz
def check_modality(filename): """ Verifica la existencia de la modalidad devuelve Falso si no se encuentra la modalidad, de lo contrario Verdadero """ end = filename.find('.nii.gz') modality = filename[end-4:end] for mod in modality: if not(ord(mod)>=48 and ord(mod)<=57): #si no está en dígitos del 0 al 9 return False return Truedef rename_for_single_modality(directory): for file in os.listdir(directory): if check_modality(file)==False: new_name = file[:file.find('.nii.gz')]+"_0000.nii.gz" os.rename(os.path.join(directory,file),os.path.join(directory,new_name)) print(f"Renombrado a {new_name}") else: print(f"Modalidad presente: {file}")rename_for_single_modality(train_image_dir)# rename_for_single_modality(test_dir)Configuración del archivo JSON
¡Casi hemos terminado!
Principalmente necesitas modificar 2 cosas:
- La modalidad (si es CT o MRI, esto cambia la normalización)
- Las etiquetas: Ingresa tus propias clases
overwrite_json_file = True #hazlo Verdadero si quieres sobrescribir el archivo dataset.json en Task_folderjson_file_exist = Falseif os.path.exists(os.path.join(task_folder_name,'dataset.json')): print('¡dataset.json ya existe!') json_file_exist = Trueif json_file_exist==False or overwrite_json_file: json_dict = OrderedDict() json_dict['name'] = task_name json_dict['description'] = "Segmentación de escaneos T1 de MindBoggle" json_dict['tensorImageSize'] = "3D" json_dict['reference'] = "ver sitio web del desafío" json_dict['licence'] = "ver sitio web del desafío" json_dict['release'] = "0.0" ######################## MODIFICA ESTO ######################## #puedes mencionar más de una modalidad json_dict['modality'] = { "0": "MRI" } #se deben mencionar labels+1 para todas las etiquetas en el conjunto de datos json_dict['labels'] = { "0": "No cerebro", "1": "Materia gris cortical", "2": "Materia blanca cortical", "3" : "Cerebelo gris", "4" : "Cerebelo blanco" } ############################################################# train_ids = os.listdir(train_label_dir) test_ids = os.listdir(test_dir) json_dict['numTraining'] = len(train_ids) json_dict['numTest'] = len(test_ids) #no hay modalidad en la imagen de entrenamiento y etiquetas en dataset.json json_dict['training'] = [{'image': "./imagesTr/%s" % i, "label": "./labelsTr/%s" % i} for i in train_ids] #eliminando la modalidad del nombre de la imagen de prueba para guardarla en dataset.json json_dict['test'] = ["./imagesTs/%s" % (i[:i.find("_0000")]+'.nii.gz') for i in test_ids] with open(os.path.join(task_folder_name,"dataset.json"), 'w') as f: json.dump(json_dict, f, indent=4, sort_keys=True) if os.path.exists(os.path.join(task_folder_name,'dataset.json')): if json_file_exist==False: print('¡dataset.json creado!') else: print('¡dataset.json sobrescrito!')Preprocesar los datos para el formato nnU-Net
Esto crea el conjunto de datos para el formato nnU-Net
# -t 1 significa "Task001", si tienes una tarea diferente ¡cámbialo!nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrityEntrenar los modelos
¡Ahora estamos listos para entrenar los modelos!
Para entrenar la U-Net 3D:
# ¡entrena la U-Net 3D a resolución completa!nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --npz Para entrenar la U-Net 2D:
# ¡entrena la U-Net 2D!nnUNet_train 2d nnUNetTrainerV2 1 0 --npzPara entrenar el modelo en cascada:
# ¡entrena la U-Net 3D en cascada!nnUNet_train 3d_lowres nnUNetTrainerV2CascadeFullRes 1 0 --npz!nnUNet_train 3d_fullres nnUNetTrainerV2CascadeFullRes 1 0 --npzNota: Si pausas el entrenamiento y deseas reanudarlo, agrega un “-c” al final para “continuar”.
Por ejemplo:
# ¡entrena la U-Net 3D a resolución completa!nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --npz Inferencia
Ahora podemos ejecutar la inferencia:
result_dir = os.path.join(task_folder_name, 'nnUNet_Prediction_Results')make_if_dont_exist(result_dir, overwrite=True)# -i es la carpeta de entrada# -o es donde deseas guardar las predicciones# -t 1 significa tarea 1, cámbialo si tienes un número de tarea diferente# Usa -m 2d, o -m 3d_fullres, o -m 3d_cascade_fullres!nnUNet_predict -i /content/drive/MyDrive/neurosciences-segmentation/my_nnunet/nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data/Task001/imagesTs -o /content/drive/MyDrive/neurosciences-segmentation/my_nnunet/nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data/Task001/nnUNet_Prediction_Results -t 1 -tr nnUNetTrainerV2 -m 2d -f 0 --num_threads_preprocessing 1Visualización de las predicciones
Primero, veamos la pérdida de entrenamiento. Esto parece muy saludable, y tenemos un Dice Score > 0.9 (curva verde).
Esto es verdaderamente excelente para tan poco trabajo y una tarea de segmentación de neuroimagen en 3D.
Echemos un vistazo a una muestra:
¡Los resultados son realmente impresionantes! Es claro que el modelo ha aprendido efectivamente cómo segmentar imágenes cerebrales con alta precisión. Si bien puede haber imperfecciones menores, es importante recordar que el campo de segmentación de imágenes avanza rápidamente y estamos dando pasos significativos hacia la perfección.
En el futuro, hay espacio para optimizar aún más el rendimiento de nnU-Net, pero eso será para otro artículo.
Si encontraste este artículo esclarecedor y beneficioso, considera seguirme para exploraciones más profundas en el mundo del aprendizaje profundo. Tu apoyo me ayuda a seguir produciendo contenido que contribuye a nuestra comprensión colectiva.
Ya sea que tengas comentarios, ideas para compartir, quieras trabajar conmigo o simplemente quieras saludar, por favor completa el formulario a continuación y comencemos una conversación.
¡Di hola! 🌿
No dudes en dejar un aplauso o seguirme para más contenido.
Referencias
- Ronneberger, O., Fischer, P., & Brox, T. (2015). U-net: Redes convolucionales para la segmentación de imágenes biomédicas. En International Conference on Medical image computing and computer-assisted intervention (pp. 234–241). Springer, Cham.
- Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: un método autoconfigurable para la segmentación de imágenes biomédicas basado en aprendizaje profundo. Nature Methods, 18(2), 203–211.
- Ioffe, S., & Szegedy, C. (2015). Normalización de lotes: acelerando el entrenamiento de redes neuronales profundas al reducir el cambio covariante interno. arXiv preprint arXiv:1502.03167.
- Ulyanov, D., Vedaldi, A., & Lempitsky, V. (2016). Normalización de instancias: el ingrediente que falta para una rápida estilización. arXiv preprint arXiv:1607.08022.
- Conjunto de datos MindBoggle
