MaskMedPaint: Medical Inpainting Techniques

• MaskMedPaint is a family of masked inpainting methods that preserve key anatomical features while mitigating spurious correlations in medical images.
• The diffusion-based approach fine-tunes pretrained models using ROI masks to adapt source images to target domain styles, enhancing classifier generalization.
• Shape-aware mask synthesis employs superpixel segmentation and organ topology priors to generate anatomically plausible masks, improving reconstruction fidelity.

MaskMedPaint refers to a family of inpainting methodologies for medical images designed to address spurious correlations and the need for anatomically-preserving, bias-mitigated data augmentation. Two principal approaches—diffusion-based masked inpainting for domain adaptation and shape-aware mask synthesis for anatomically plausible training—have emerged under this term. These approaches leverage distinct strategies but share the goal of improving generalization in image classification and reconstruction, particularly in high-stakes clinical domains where spurious features can result in vulnerable or unreliable models (Jin et al., 2024, Yeganeh et al., 2022).

1. Background and Motivation

Medical image analysis models, especially classifiers, are prone to exploiting non-causal, dataset-specific shortcuts (“spurious features”) such as ruler markings or radiology markers that are inadvertently correlated with pathology labels in the training set. This exploitation leads to sharp performance drops when evaluating models in alternate domains (e.g., new hospitals, different scanner vendors), thereby undermining robustness and clinical trust. Black-box model adjustments are often impractical since clinicians prefer data-driven interventions that can be visually inspected and vetted.

Diffusion-based generative augmentation—effective in natural images—faces limitations in medical contexts due to the difficulty of identifying and describing subtle, domain-specific spurious features via natural language prompts. MaskMedPaint addresses this by (i) spatially constraining generative augmentation to preserve the lesion or anatomical region of interest (ROI), and (ii) shifting only the surrounding context to match the style statistics of the target domain (Jin et al., 2024).

In a separate but related line, shape-aware mask generation for inpainting aims to learn statistical priors on organ topology, overcoming the artifacts and trivial hole-filling induced by random or geometric masking strategies. Here, mask generation itself is tailored to reflect the shape diversity of anatomical structures (Yeganeh et al., 2022).

2. Pipeline and Methodological Overview

2.1. Diffusion-based Masked Inpainting for Domain Adaptation

MaskMedPaint’s main pipeline comprises three sequential stages:

• Class-Conditional Source Fine-Tuning: Starting from a pretrained diffusion model (Stable Diffusion 1.5), fine-tune on labeled source images, appending each sample’s class label (e.g., “malignant”) to the prompt. The model learns associations between class tokens and characteristic ROI (e.g., lesion) appearances.
• Unconditional Target Fine-Tuning (DreamBooth): Remove the ROI from a small set of unlabeled target images using an inpainting model (LaMa), producing “background-only” samples. DreamBooth-style finetuning then associates a new token (e.g., “target-domain”) with the style and artifact statistics of these target backgrounds, without class information.
• Masked Inpainting Generation: For each labeled source image, compute a binary ROI mask $M$ (from a U-Net segmenter). During diffusion sampling, ROI voxels are preserved while the context is denoised toward the target-domain style through masked guidance, maintaining lesion integrity but shifting the background.

The composite partial image at diffusion time step $t$ takes the form:

$x'_t = M \odot x_t + (1 - M) \odot \tilde{x}_t$

where $x_t$ is a forward-diffused source and $\tilde{x}_t$ is the inpainted background conditioned on the class and domain tokens.

After generating a sufficiently large augmented dataset, a downstream classifier (DenseNet121) is trained using both original and MaskMedPaint-augmented images (Jin et al., 2024).

2.2. Shape-Aware Mask Generation for Inpainting

In an alternative MaskMedPaint approach, mask sampling is designed to better align with organ topology, avoiding the training pitfalls of square or arbitrary irregular masking. The procedure consists of:

• Superpixel Decomposition: Apply Felzenszwalb’s over-segmentation to the input, producing pseudo-segments that reflect local anatomical structure.
• Weight-Biased Segment Sampling: Sample segments with a bias toward smaller (harder-to-reconstruct) patches and optionally merge neighboring segments for diversity.
• Spatial Placement Strategies: Position sampled segments according to three strategies—random, inside a morphologically estimated organ region (IBR), or at their original position (OPS)—to maximize anatomical plausibility during inpainting training.
• Rule-Based Mask Generator: The sampler is non-trainable, leveraging only image graph statistics and simple morphological operations (Yeganeh et al., 2022).

3. Diffusion Model and Inpainting Formalism

MaskMedPaint is built on the Denoising Diffusion Probabilistic Model (DDPM) framework. The forward process gradually corrupts images with Gaussian noise:

$$q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t}\,x_{t-1}, \beta_t I)$$

and the reverse process iteratively denoises in the U-Net backbone:

$$p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))$$

In MaskMedPaint, at each reverse denoising step, the ROI mask $M$ is used to prevent diffusion updates within the lesion, ensuring only the contextual background is synthesized toward the target domain’s statistics.

For the shape-aware masking pipeline, inpainting is driven by CTSDG’s dual-branch U-Net, with losses computed for both pixel intensities and structural edges, benefiting from the anatomically plausible mask sampling protocol (Yeganeh et al., 2022).

4. Empirical Evaluation and Results

4.1. Clinical Datasets and Metrics

Both MaskMedPaint pipelines were tested on standard benchmarks:

• ISIC 2018 Dermoscopy: Targeting spurious correlation between ruler markings and malignant labels. The base model achieves 0.93 accuracy (source), but only 0.146 (target domain). MaskMedPaint achieves 0.746 (source), 0.344 (target), closely approaching the oracle classifier trained on labeled target data (0.742 target). MaskMedPaint correctly flips background artifacts in 40–47% of images (Jin et al., 2024).
• Chest X-rays: Transferring from MIMIC-CXR to NIH ChestXray14, target AUROC increases from 0.490 (base) to 0.546 (MaskMedPaint), with competitive source performance.
• Natural Image Transfer: On Waterbirds and iWildCam (grayscale-to-color, night-to-day), MaskMedPaint improves target accuracy (e.g., Waterbirds: 0.264→0.571).

In shape-aware masking, abdominal MRI organ inpainting on the CHAOS dataset demonstrates that pseudo-segment (PS) masks yield higher PSNR and SSIM than square or irregular masks (e.g., PSNR: 19.38 vs. 18.07 for squares). Positioning masks using organ-region priors (IBR, OPS) further improves reconstruction fidelity, with all pairwise statistical tests (square vs. PS; random vs. IBR/OPS) yielding $p<0.01$ (Yeganeh et al., 2022).

4.2. Ablations and Analysis

• Quantity of Augmented Data: 2,500 MaskMedPaint samples match the performance boost of 100 real target samples; returns saturate past 5,000 augmentations.
• Method Comparison: Pure text-to-image or image-to-image editing is ineffective in medical settings, failing to preserve critical structure or adequately remove spurious context; ROI masking is essential.
• Shape-Awareness: Mask shape and position that align with organ topology yield significant PSNR/SSIM gains, confirming the importance of anatomically-informed masking for medical inpainting (Yeganeh et al., 2022).

5. Implementation Details

Key implementation characteristics include:

• Diffusion Backbones: Stable Diffusion v1.5 with text conditioning; DreamBooth for domain-specific background adaptation.
• ROI Segmentation: U-Net or other medical segmenters for binary ROI mask generation; in DreamBooth, ROI is excised with LaMa inpainting.
• Classifier: DenseNet121 pretrained on ImageNet, fine-tuned on the augmented dataset.
• Training and Preprocessing: All images are resized to $224\times224$ and normalized; mask generation or pseudo-segment extraction employs scikit-image’s implementation with specified thresholding parameters.
• Inpainting Networks: For shape-aware masking, a dual-branch U-Net trained with adversarial, perceptual, style, and edge consistency losses. The mask generator leverages morphological opening (e.g., $t$0 structuring element) for organ-region detection.

6. Limitations and Future Directions

MaskMedPaint pipelines exhibit key limitations:

• Dependence on high-quality ROI segmentation; errors can result in destruction or duplication of spurious cues.
• Requirement for approximately 100 unlabeled target-domain images to prevent overfitting during DreamBooth adaptation.
• Augmentation is not a substitute for real target data in complex settings, particularly with sensitive-attribute or intersectional bias.

Suggested future research includes automated spurious feature discovery, controlled spatial artifact generation, and integration with model-based robustness methods such as Distributionally Robust Optimization or invariant risk minimization (Jin et al., 2024). For shape-aware masking, extension to different anatomies and modalities is feasible, provided an organ-bounding region can be estimated.

7. Significance and Position within the Literature

MaskMedPaint introduces two central advances: (1) a clinically auditable, domain-adaptive inpainting pipeline capable of mitigating shortcut learning in medical classifiers without reliance on precise textual descriptions of spurious features, and (2) a shape-aware, unsupervised mask generation method that injects anatomical plausibility into inpainting networks, improving the fidelity of reconstructions and statistical power of empirical evaluations.

Both variants are designed to be compatible with established generative and discriminative architectures, require no additional labeled data, and demonstrate substantial performance gains across modalities and datasets. The integration of visual auditability, domain-specific adaptation, and anatomical awareness distinguishes MaskMedPaint from prior text-to-image and generic mask-based augmentation approaches (Jin et al., 2024, Yeganeh et al., 2022).