Diffusion Model Based Brain Tumour Segmentation Enhanced With Inpainting Method

Resumen

En este trabajo se presenta un método para la segmentación de tumores cerebrales en imágenes MRI basado en modelos de difusión. La propuesta se centra en un enfoque débilmente supervisado, que evita la necesidad de etiquetas detalladas a nivel de píxel, lo que lo hace especialmente útil en entornos donde la anotación manual de datos es costosa y limitada.

El proyecto consta de dos fases principales. En la primera, se entrena un modelo de difusión DDPM con imágenes de cerebros sanos, permitiendo generar versiones reconstruidas de imágenes con tumores, en las cuales las regiones patológicas son reemplazadas por estructuras saludables. A partir de la diferencia entre la imagen original y la reconstruida, se obtiene una máscara de anomalía que destaca las áreas afectadas. Para refinar esta máscara, se aplica el método de restauración de imagen RePaint, lo que mejora la calidad de la segmentación final.

En la segunda parte del estudio, se replica un modelo DDIM guiado por un clasificador, tomado de la literatura previa, con el objetivo de comparar su desempeño con el método propuesto. Este modelo genera una versión sana de una imagen con tumor mediante condicionamiento del clasificador y, posteriormente, se aplica el mismo proceso de RePaint para obtener la segmentación final.

Los resultados obtenidos muestran que la combinación de modelos de difusión con técnicas de inpainting permite mejorar la calidad de la segmentación, proporcionando un método alternativo y eficiente para la detección de tumores cerebrales. Además, este enfoque reduce la dependencia de datos anotados.

Abstract:

In this work, an method for brain tumour segmentation in MRI images based on diffusion models is presented. The proposed approach focuses on a weakly supervised method, which eliminates the need for detailed pixel-level annotations, making it particularly useful in environments where manual data labelling is costly and limited.

The project consists of two main phases. In the first phase, a diffusion model DDPM is trained using images of healthy brains, enabling the generation of reconstructed versions of tumour-affected images, where pathological regions are replaced with healthy structures. The anomaly mask is then obtained by computing the difference between the original and reconstructed images, highlighting the affected areas. To refine this mask, the inpainting method RePaint is applied, enhancing the quality of the final segmentation.

In the second phase of the study, a classifier-guided DDIM model, taken from previous literature, is replicated to compare its performance with the proposed method. This model generates a healthy version of a tumour-affected image using classifier conditioning and subsequently applies the same RePaint process to obtain the final segmentation.

The results obtained demonstrate that combining diffusion models with inpainting techniques improves segmentation quality, providing an alternative and efficient method for brain tumour detection. Furthermore, this approach reduces the reliance on annotated data.