IterMask2: Iterative Unsupervised Anomaly Segmentation via Spatial and Frequency Masking for Brain Lesions in MRI

TL;DR

The paper tackles unsupervised anomaly segmentation in brain MRI by learning only from healthy tissue and detecting pathologies as deviations from the learned normal distribution. It introduces IterMask2, which combines iterative spatial mask refinement during inference with Fourier-derived high-frequency masking to guide reconstruction of masked regions, aided by a UNet trained on randomly masked healthy data. The approach achieves state-of-the-art or near-state-of-the-art performance on BraTS, ATLAS, and ISLES datasets, with reconstruction quality on normal tissue reaching SSIM values around $0.985$--$0.997$ across modalities and Dice scores that exceed several baselines. The work demonstrates the value of jointly refining masks and leveraging high-frequency structural cues, offering robust performance and potential for interactive clinician-in-the-loop use and extension to 3D data.

Abstract

Unsupervised anomaly segmentation approaches to pathology segmentation train a model on images of healthy subjects, that they define as the 'normal' data distribution. At inference, they aim to segment any pathologies in new images as 'anomalies', as they exhibit patterns that deviate from those in 'normal' training data. Prevailing methods follow the 'corrupt-and-reconstruct' paradigm. They intentionally corrupt an input image, reconstruct it to follow the learned 'normal' distribution, and subsequently segment anomalies based on reconstruction error. Corrupting an input image, however, inevitably leads to suboptimal reconstruction even of normal regions, causing false positives. To alleviate this, we propose a novel iterative spatial mask-refining strategy IterMask2. We iteratively mask areas of the image, reconstruct them, and update the mask based on reconstruction error. This iterative process progressively adds information about areas that are confidently normal as per the model. The increasing content guides reconstruction of nearby masked areas, improving reconstruction of normal tissue under these areas, reducing false positives. We also use high-frequency image content as an auxiliary input to provide additional structural information for masked areas. This further improves reconstruction error of normal in comparison to anomalous areas, facilitating segmentation of the latter. We conduct experiments on several brain lesion datasets and demonstrate effectiveness of our method. Code is available at: https://github.com/ZiyunLiang/IterMask2

Figures (4)

• Figure 1: Overview of proposed method. (a) The frequency masking strategy that provides structural information for an image via its high frequency components. (b) Training of the reconstruction model. Random spatial masking is applied to input image $\textbf{x}$. High image frequencies from step (a) are given as auxiliary input $\textbf{x}_c$ to the model. (c) shows the iterative mask refinement process during inference. The mask $\textbf{m}_t$ of iteration t gradually shrinks toward the anomaly with the guidance of spatially unmasked area $\textbf{x}_m$ and high image frequencies $\textbf{x}_c$.
• Figure 2: Performance on ISLES and ATLAS datasets. Best marked in bold.
• Figure 2: Visualization result for FLAIR(BraTS).
• Figure 3: Sensitivity Analysis.