Diffusion Models for Counterfactual Generation and Anomaly Detection in Brain Images

Segmentation masks of pathological areasare useful in many medical applications, such as braintumour and stroke management. Moreover, healthy counterfactualsof diseased images can be used to enhance radiologists’training files and to improve the interpretabilityof segmentation models. In this work, we present a weaklysupervised method to generate a healthy version of a diseasedimage and then use it to obtain a pixel-wise anomalymap. To do so, we start by considering a saliency map thatapproximately covers the pathological areas, obtained withACAT. Then, we propose a technique that allows to performtargeted modifications to these regions, while preservingthe rest of the image. In particular, we employ a diffusionmodel trained on healthy samples and combine DenoisingDiffusion Probabilistic Model (DDPM) and Denoising DiffusionImplicit Model (DDIM) at each step of the samplingprocess. DDPM is used to modify the areas affected bya lesion within the saliency map, while DDIM guaranteesreconstruction of the normal anatomy outside of it. Thetwo parts are also fused at each timestep, to guarantee thegeneration of a sample with a coherent appearance anda seamless transition between edited and unedited parts.We verify that when our method is applied to healthy samples,the input images are reconstructed without significantmodifications. We compare our approach with alternativeweakly supervised methods on the task of brain lesionsegmentation, achieving the highest mean Dice and IoUscores among the models considered.