Selective Masking based Self-Supervised Learning for Image Semantic Segmentation

TL;DR

Problem: Masked image modeling pretraining often relies on random masking and large budgets, limiting efficiency for semantic segmentation in resource-constrained settings. Approach: Introduce Selective Masking Image Reconstruction with iterative partitioned pretraining and loss-driven patch masking on a CNN/CNN-based U-Net backbone, using MS-SSIM+L1 losses. Findings: Outperforms random masking and ImageNet pretraining across Pascal VOC, Cityscapes, and weed datasets (Nassar, Sugarbeets), with notable gains on underrepresented classes; self-pretraining typically surpasses cross-pretraining. Impact: Provides a practical, data-efficient SSL pretraining option for end-to-end segmentation in low-resource or domain-specific scenarios, reducing reliance on large-scale pretraining.

Abstract

This paper proposes a novel self-supervised learning method for semantic segmentation using selective masking image reconstruction as the pretraining task. Our proposed method replaces the random masking augmentation used in most masked image modelling pretraining methods. The proposed selective masking method selectively masks image patches with the highest reconstruction loss by breaking the image reconstruction pretraining into iterative steps to leverage the trained model's knowledge. We show on two general datasets (Pascal VOC and Cityscapes) and two weed segmentation datasets (Nassar 2020 and Sugarbeets 2016) that our proposed selective masking method outperforms the traditional random masking method and supervised ImageNet pretraining on downstream segmentation accuracy by 2.9% for general datasets and 2.5% for weed segmentation datasets. Furthermore, we found that our selective masking method significantly improves accuracy for the lowest-performing classes. Lastly, we show that using the same pretraining and downstream dataset yields the best result for low-budget self-supervised pretraining. Our proposed Selective Masking Image Reconstruction method provides an effective and practical solution to improve end-to-end semantic segmentation workflows, especially for scenarios that require limited model capacity to meet inference speed and computational resource requirements.

Figures (5)

• Figure 1: Our proposed Selective Masking Image Reconstruction SSL method iteratively trains on ten equally divided partitions of the pretraining dataset. The model is initialized by training on randomly masked partition 0. We then selectively mask partition $i$ using the model trained on partition $i-1$. Model weights from partition $i-1$ training are used to initialize the model for partition $i$, which is then trained on selectively masked partition $i$. We repeat this selective masking and training process until all partitions are used.
• Figure 2: Our proposed selective masking process starts by generating five 50% randomly masked samples, which are then reconstructed by the trained U-Net model. A patch loss map is calculated for each reconstructed sample, and then the sample patch loss maps are aggregated and thresholded (top 50%) to generate a selective mask. The selective mask is applied to the original image, resulting in the desired selectively masked image.
• Figure 3: Examples comparing randomly and selectively masked images. Left-to-right: two examples each for Pascal VOC, Cityscapes, Nassar2020, and Sugarbeets 2016. The brightness of Sugarbeets images is increased to improve visibility. We can observe selective masking following the edges of leaves and roads in examples in the second column. Other examples show that selective masks are generally concentrated on non-background objects.
• Figure 4: Selectively masked and reconstructed partition 9 examples from PascalVOC (a), Cityscapes (b), Nassar 2020 (c), and Sugarbeets 2016 (d). The brightness of Sugarbeets images is increased to improve visibility. The images in each column from left to right are: the selectively masked image, the reconstructed image, and the original image. Examples are randomly cropped to 256$\times$256 for validation, resulting in varied masked ratios. The Pascal VOC (a) and Cityscapes (b) examples show the selective masking process actively masks the entirety or majority of non-background object patches.
• Figure 5: Comparison of selective masking and random masking reconstruction examples from partition 9 showing the original image (a), randomly masked and reconstructed images (b), and selectively masked and reconstructed images (c). The examples from the top to the bottom row are from Pascal VOC, Cityscapes, Nassar 2020, Sugarbeets 2016. The brightness of Sugarbeets images is increased to improve visibility.