This repository implements scale-aware self-supervised learning (SSL) for targeted segmentation of small and sparse structures, as proposed in Scale-Aware Self-Supervised Learning for Segmentation of Small and Sparse Structures (ICASSP 2026). It integrates small-window cropping into standard SSL pipelines (SimCLR, BYOL, VICReg) to bias representations toward fine-scale features. The method is evaluated on seismic fault segmentation and neuroimaging cell/vessel segmentation under label-constrained regimes, showing consistent gains for small targets while leaving large-scale segmentation largely unchanged.
You can download the datasets used on this benchamrk in these links: CRACKS MTNeuro
Make sure the following Python packages are installed:
pip install -r requirements.txt
Train a model: Each model script is controlled via a YAML config file. Example:
main_patch_lightly.py --config config/config.yaml
You can customize training strategy, dataset used and model hyperparameters either by using the config file or overriding the corresponding flag using the command line. Training logs are printed and optionally visualized via Weights & Biases.
Run inference: In order to run inference using a saved model run:
python segmentation_train.py --config config/config-seg.yaml