ReDiMask

Reference implementation of

ReDiMask: Regionized Diffusion Masking for Annotation-Free Remote Sensing Building Extraction

ReDiMask converts frozen Stable Diffusion priors into reliable building supervision without any pixel-level annotation. Diffusion groups are informative building candidates but unreliable masks: they mix buildings with visually similar non-building regions, and building pixels are scattered across several groups. ReDiMask separates reliable supervision extraction from the recovery of building regions omitted by that supervision:

Stage	Module	Role
1	RDL — Regionized Diffusion Labeling	Turns frozen diffusion features into conservative pseudo masks (reliable building / reliable non-building / ignored)
2	Coarse student	Learns a dense building prediction from the valid pseudo-label pixels
3	W-MR — Warm-start Mask Refinement	Recovers omitted building regions with gated positive residual updates anchored on the coarse prediction
4	ICR — Interior Consistency Correction	Fills small reliable interior holes at inference time

Ground truth is used for evaluation and for the Fig. 1 mismatch analysis only. It never enters pseudo-label construction, coarse-student training, W-MR training or test-time mask generation.

Installation

python -m venv .venv
source .venv/bin/activate          # Windows: .venv\Scripts\activate
pip install -r requirements.txt

Stable Diffusion v1-5 is loaded through diffusers. Either set diffusion.model_path to a local diffusers directory, or leave the default runwayml/stable-diffusion-v1-5 and make sure the weights are reachable.

Data layout

Each benchmark is prepared as non-overlapping 256x256 patches. A manifest is a CSV with one row per patch:

image,stem
/data/vaihingen/train/images/area1_p1.png,area1_p1

Label columns (target, mask, mask_path, ...) are recognised but are not required except for evaluation and the diagnostic analysis.

Usage

All tools read the same YAML configuration, which already contains the paper settings (configs/redimask.yaml).

CFG=configs/redimask.yaml

# 1. Cache frozen Stable Diffusion features at T = {25, 50, 150}
python tools/extract_features.py \
    --manifest  data/vaihingen_train.csv \
    --cache-dir work/features/vaihingen \
    --config    $CFG

# 2. Build conservative RDL pseudo masks (also fits the diffusion groups)
python tools/build_pseudo_labels.py \
    --manifest  data/vaihingen_train.csv \
    --cache-dir work/features/vaihingen \
    --out-dir   work/rdl/vaihingen \
    --config    $CFG

# 3. Train the coarse segmentation student
python tools/train_coarse_student.py \
    --manifest work/rdl/vaihingen/pseudo_label_manifest.csv \
    --out-dir  work/coarse/vaihingen \
    --config   $CFG

# 4. Train W-MR on top of the frozen coarse student
python tools/train_wmr.py \
    --manifest         work/rdl/vaihingen/pseudo_label_manifest.csv \
    --coarse-checkpoint work/coarse/vaihingen/best_pseudo_ce_model.pth \
    --out-dir          work/wmr/vaihingen \
    --config           $CFG

# 5. Inference + evaluation
python tools/infer.py \
    --manifest          data/vaihingen_test.csv \
    --coarse-checkpoint work/coarse/vaihingen/best_pseudo_ce_model.pth \
    --wmr-checkpoint    work/wmr/vaihingen/best_pseudo_ce_model.pth \
    --gt-dir            /data/vaihingen/test/masks \
    --out-dir           work/eval/vaihingen \
    --config            $CFG

Evaluate an existing prediction directory independently:

python tools/evaluate.py \
    --pred-dir work/eval/vaihingen/pred_masks \
    --manifest data/vaihingen_test.csv \
    --gt-dir   /data/vaihingen/test/masks \
    --out-dir  work/metrics/vaihingen \
    --config   $CFG

Reproduce the diffusion-group mismatch analysis of Fig. 1 by adding --dump-group-maps to step 2 and running:

python tools/diagnose_groups.py \
    --manifest  data/vaihingen_train.csv \
    --group-dir work/rdl/vaihingen/group_maps \
    --gt-dir    /data/vaihingen/train/masks \
    --out-dir   work/analysis/vaihingen \
    --config    $CFG

Ground-truth encoding

Benchmark ground truth is not encoded uniformly: the converted patches use 0/1, binary releases such as Waterloo use 0/255, and semantic labels use arbitrary ids. Both infer.py and evaluate.py therefore accept

--gt-building-value 255      # value that marks building pixels in the ground truth
--building-value 1           # value that marks building pixels in the prediction

Both default to data.building_id.

Tests

python tests/smoke_test.py        # unit-level checks of every module
python tests/end_to_end_test.py   # runs every pipeline stage on synthetic data

Neither test downloads Stable Diffusion weights or a benchmark dataset.

Configuration

Every hyper-parameter reported in the paper lives in configs/redimask.yaml:

Key	Value	Paper
diffusion.timesteps	[25, 50, 150]	low-noise timestep set T
diffusion.hook_module	mid_block	U-Net mid-block activation, 16x16x1280
region.n_segments / compactness	50 / 8.0	SLIC partition
rdl.num_groups	6	K diffusion groups
rdl.num_candidate_groups	2	Kc retained groups
rdl.reference_timestep	50	reference for cross-timestep alignment
rdl.weights	(0.4, 0.3, 0.2, 0.1)	(lambda_t, lambda_m, lambda_s, lambda_r)
rdl.min_support_ratio / max_support_ratio	0.005 / 0.70	group support filter
rdl.positive_votes	2	tau+
wmr.tau_c / tau_l / tau_h	0.5 / 0.3 / 0.7	coarse threshold and uncertainty band
wmr.boundary_radius	3	boundary band radius r
wmr.infer_steps	3	Q fixed refinement steps
wmr.alpha	1.0	residual scale in Eq. (21); not fixed by the paper
wmr.noise_train	[0.0, 0.30]	perturbation range in Eq. (17); not fixed by the paper
icr.max_area / min_prob	1024 / 0.15	ICR conditions
data.boundary_tolerance	3	BoundF matching tolerance

Metrics

IoU, F1 and BoundF are reported for the building class. BoundF extracts building boundaries from binary masks and matches predicted and ground-truth boundary pixels within a 3-pixel dilation tolerance. Building Ratio (BR) and Recall Contribution (RC) are available through redimask.metrics for diagnostic pixel sets such as newly added foreground pixels.

Repository layout

redimask/
  features.py    frozen SD v1-5 multi-timestep mid-block features
  regions.py     SLIC partition and region descriptors (Eq. 1)
  rdl.py         Regionized Diffusion Labeling (Eqs. 1-12)
  student.py     coarse segmentation student (Eqs. 13-16)
  wmr.py         Warm-start Mask Refinement (Eqs. 17-25)
  icr.py         Interior Consistency Correction
  metrics.py     IoU / F1 / BoundF / BR / RC
  analysis.py    Fig. 1 mismatch analysis
  data.py        manifest and dataset helpers
  nets.py        shared U-Net backbone
tools/           command line entry points for every stage
configs/         paper configuration
docs/            equation-to-code mapping

See docs/paper_code_mapping.md for the mapping between the paper equations and the implementation, including the few places where the paper leaves a functional form unspecified.

Citation

The full citation will be added upon publication.

License

Released under the MIT License.