FREST: Feature RESToration for Semantic Segmentation under Multiple Adverse Conditions

[ECCV 2024] Official PyTorch implementation of FREST.

FREST is a source-free domain adaptation (SFDA) framework for semantic segmentation under adverse conditions. It alternates two steps: (1) learning a condition embedding space that isolates condition-specific information, and (2) restoring features of adverse-condition images toward the normal condition on that space. This repository covers the main Cityscapes → ACDC setting.

Method overview

Paper component	Where in the code
Condition strainer psi_strainer (Fig. 3)	models/backbones/strainer.py (ConditionStrainer), attached per stage in models/backbones/mix_transformer_strainer.py
Projection onto the condition embedding space psi_proj	models/heads/projection.py (ProjectionHead) + models/backbones/condition_projector.py (ConditionProjector)
Condition-specific loss L_spec (Eq. 1–2)	sampling in models/losses.py (ConditionContrastiveSampler), contrastive term in FREST.training_step (Step 1)
Feature restoration loss L_resto (Eq. 3)	restoration_loss in FREST.training_step (Step 2)
Adverse-condition discriminating loss L_dis (Eq. 4)	discriminate=True path in the backbone; discriminating_loss in FREST.training_step
Self-training L_self (CBST) + entropy L_ent	helpers/pseudo_labels.py, models/losses.py (NormalizedEntropyLoss)
Restored features used at inference (encoder + decoder only)	inference=True path in the backbone (the strainer is skipped)

The two steps are alternated every iteration; only the encoder and decoder are used for evaluation.

Setup

conda create -n frest python=3.8 -y && conda activate frest
pip install -r requirements.txt

A CUDA toolchain is required: the optical-flow correlation op (models/correlation_ops) is JIT-compiled on first use (needs nvcc + ninja).

Data and weights

Set DATA_DIR to a folder laid out as:

$DATA_DIR/
  ACDC/                 # rgb_anon/, gt/  (https://acdc.vision.ee.ethz.ch/)
  pseudo_labels/
    pseudo_labels_train_ACDC_cma_segformer/   # auto-downloaded on first run

• Source model: place the Cityscapes-pretrained SegFormer-B5 at ./pretrained_models/segformer.b5.1024x1024.city.160k.pth (official SegFormer weights). download link
• Alignment (training only): place the UAWarpC MegaDepth weights at ./pretrained_models/uawarpc_megadepth.ckpt. Not needed for evaluation. download link

See pretrained_models/README.md for download links of the two files above.

Train

DATA_DIR=/path/to/data bash scripts/train_acdc.sh

Evaluate

DATA_DIR=/path/to/data bash scripts/eval_acdc.sh checkpoints/frest_acdc.ckpt

FREST checkpoint

Download the FREST checkpoint and place it at checkpoints/frest_acdc.ckpt:

• Checkpoint (Cityscapes → ACDC): download link

If you instead have a checkpoint from the original training code (legacy module names), convert its keys to this release's layout with:

python convert_checkpoint.py SRC.ckpt checkpoints/frest_acdc.ckpt

Citation

@inproceedings{lee2024frest,
  title     = {FREST: Feature RESToration for Semantic Segmentation under Multiple Adverse Conditions},
  author    = {Lee, Sohyun and Kim, Namyup and Kim, Sungyeon and Kwak, Suha},
  booktitle = {European Conference on Computer Vision (ECCV)},
  year      = {2024}
}

Acknowledgements

The codebase builds on Refign/CMA and SegFormer.