Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Diffusion Models for Counterfactual Generation and Anomaly Detection in Brain Images

Alessandro Fontanella, Grant Mair, Joanna Wardlaw, Emanuele Trucco, Amos Storkey

TL;DR

This work presents a weakly supervised method to generate a healthy version of a diseased image and then uses it to obtain a pixel-wise anomaly map and compares its approach with alternative weakly supervised methods on the task of brain lesion segmentation.

Abstract

Segmentation masks of pathological areas are useful in many medical applications, such as brain tumour and stroke management. Moreover, healthy counterfactuals of diseased images can be used to enhance radiologists' training files and to improve the interpretability of segmentation models. In this work, we present a weakly supervised method to generate a healthy version of a diseased image and then use it to obtain a pixel-wise anomaly map. To do so, we start by considering a saliency map that approximately covers the pathological areas, obtained with ACAT. Then, we propose a technique that allows to perform targeted modifications to these regions, while preserving the rest of the image. In particular, we employ a diffusion model trained on healthy samples and combine Denoising Diffusion Probabilistic Model (DDPM) and Denoising Diffusion Implicit Model (DDIM) at each step of the sampling process. DDPM is used to modify the areas affected by a lesion within the saliency map, while DDIM guarantees reconstruction of the normal anatomy outside of it. The two parts are also fused at each timestep, to guarantee the generation of a sample with a coherent appearance and a seamless transition between edited and unedited parts. We verify that when our method is applied to healthy samples, the input images are reconstructed without significant modifications. We compare our approach with alternative weakly supervised methods on the task of brain lesion segmentation, achieving the highest mean Dice and IoU scores among the models considered.

Diffusion Models for Counterfactual Generation and Anomaly Detection in Brain Images

TL;DR

This work presents a weakly supervised method to generate a healthy version of a diseased image and then uses it to obtain a pixel-wise anomaly map and compares its approach with alternative weakly supervised methods on the task of brain lesion segmentation.

Abstract

Segmentation masks of pathological areas are useful in many medical applications, such as brain tumour and stroke management. Moreover, healthy counterfactuals of diseased images can be used to enhance radiologists' training files and to improve the interpretability of segmentation models. In this work, we present a weakly supervised method to generate a healthy version of a diseased image and then use it to obtain a pixel-wise anomaly map. To do so, we start by considering a saliency map that approximately covers the pathological areas, obtained with ACAT. Then, we propose a technique that allows to perform targeted modifications to these regions, while preserving the rest of the image. In particular, we employ a diffusion model trained on healthy samples and combine Denoising Diffusion Probabilistic Model (DDPM) and Denoising Diffusion Implicit Model (DDIM) at each step of the sampling process. DDPM is used to modify the areas affected by a lesion within the saliency map, while DDIM guarantees reconstruction of the normal anatomy outside of it. The two parts are also fused at each timestep, to guarantee the generation of a sample with a coherent appearance and a seamless transition between edited and unedited parts. We verify that when our method is applied to healthy samples, the input images are reconstructed without significant modifications. We compare our approach with alternative weakly supervised methods on the task of brain lesion segmentation, achieving the highest mean Dice and IoU scores among the models considered.
Paper Structure (17 sections, 10 equations, 7 figures, 1 table, 1 algorithm)

This paper contains 17 sections, 10 equations, 7 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Saliency Maps
  4. Counterfactual Explanations
  5. Anomaly Detection
  6. Methods
  7. Diffusion Models
  8. Dif-fuse
  9. Training details
  10. Experiments
  11. Data
  12. Experimental Setup
  13. Counterfactual Examples
  14. Hyperparameters
  15. Quantitative Evaluation
  16. ...and 2 more sections

Figures (7)

  • Figure 1: Our approach begins by transforming an abnormal image $\bm{x}_0$ into its noised version $\bm{x}_K$ using the reversed sampling technique of DDIMs. Subsequently, we employ DDPM sampling to modify the pathological area, identified through the saliency map generated with ACAT, aiming to restore the normal anatomical structure based on the contextual information from the surrounding regions. Meanwhile, the regions of the image that do not contain any pathological elements are restored to their original appearance using DDIM sampling. Throughout the sampling process, these two components are fused together to ensure a seamless and realistic transition between the edited and unedited parts, resulting in a final image $\bm{\hat{x}}_0$ with a visually coherent and natural appearance
  • Figure 2: Original image from IST-3 (a) and healthy counterfactuals (second row) with corresponding anomaly maps (bottom row), obtained with DenoisingAE (b), f-AnoGAN (c), AnoDDPM (d), AutoDDPM (e), classifier guidance (f), classifier-free guidance (g), ACAT (h) and Dif-fuse (i). ACAT generates a reasonable anomaly map, but is not able to fully remove the lesion. Dif-fuse refines the anomaly map obtained with ACAT, while at the same time creating a credible counterfactual example. The other approaches introduce artifacts and/or identify the pathological area less correctly.
  • Figure 3: Input image from IST-3 (a) and normal image generated by applying the mask only at the end of the sampling process (b). We can observe that (b) presents some artifacts and does not have a smooth transition between edited and unedited parts.
  • Figure 4: Original image from BraTS 2021 with ground truth segmentation mask (a) and healthy images (top row) with corresponding anomaly maps (bottom row), obtained with DenoisingAE (b), f-AnoGan (c), AnoDDPM (d), AutoDDPM (e), classifier guidance (f), classifier-free guidance (g) and with Dif-fuse (h). f-Ano GAN falls short in generating believable counterfactuals, whereas the other approaches yield higher-quality results. However, DenoisingAE, AnoDDPM and AutoDDPM do not fully remove the lesion, while the counterfactuals generated with CG and CFG exhibit some artifacts.
  • Figure 5: Dice scores obtained on the validation dataset with different combinations of thresholding percentiles to binarise the saliency maps and noise amounts K. We obtain the best result with $K=500$ and pixels in the $90^{th}$ percentile of the saliency maps.
  • ...and 2 more figures