PhenoProto-SSL

Self-supervised pretraining + prototype-based finetuning for crop-type semantic segmentation on satellite image time series, built on top of a from-scratch reproduction of STCLN (XiaoleiQinn/STCLN) on the PASTIS benchmark.

Status: work in progress, actively improving. Track A (S3M pretrain, 100 epochs) and the official-baseline reproduction (Rung 0) have both completed. A full PhenoProto-SSL finetune sourced from our own S3M pretrain (s3m_ep89) beat the official baseline on the held-out fold-4 test set (mIoU 0.4298 vs 0.3493) — see Sample predictions for real classification maps and Results for the full numbers. A second run from the final pretrain checkpoint (s3m_ep99) is in progress. See Status for exactly what has and hasn't run.

---

Table of contents

• What this is
• Method
• Dataset
• Experimental protocol
• Configuration
• Hardware & environment
• Status
• Results
• Sample predictions
• Reproducing
• Repository structure
• What to push to a public repo
• Known issues & engineering notes
• Acknowledgements

---

What this is

PASTIS is a Sentinel-2 optical satellite image time series dataset for panoptic/semantic agricultural parcel segmentation — 2,433 patches of 128×128 px, ~20–60 irregularly-spaced acquisitions per patch, 10 spectral bands, 20 semantic crop-type classes (+ background/void). STCLN is a published transformer-based baseline (UTAE backbone + spatio-temporal attention fusion) for this task.

This project has two goals, run as two parallel overnight tracks on one GPU:

• Track A — pretrain the UTAE encoder with a novel self-supervised masking objective (S3M, below) on unlabeled data, before any labels are used.

• Track B — an ablation ladder that starts from the officially released pretrained encoder (not our own Track A pretrain, so the two tracks can run independently) and adds one PhenoProto-SSL component at a time, to isolate which addition actually moves the needle:

  Rung	Config	Isolates
  0	official STCLN code, unmodified	the reference number
  A	our pipeline, plain linear head	pipeline-equivalence check
  B	+ prototype head	contribution of PA-Seg-style prototypes
  C	+ balanced SupCon	contribution of contrastive geometry
  D	+ semi-supervised mean-teacher	contribution of unlabeled data

Method

Backbone (unchanged from official STCLN, imported verbatim — never reimplemented, to eliminate an entire class of silent reproduction bugs): U-TAE encoder (encoder_widths=[32,256], 2-level) → LTAE2d temporal self-attention → a spatio-temporal attention (STA) fusion block (temporal max-pool branch × spatial-attention branch, ReZero-gated).

New contributions on top of that backbone:

• S3M — Spectro-Spatio-Temporal Masking (masking.py). The official S3M baseline masks a spatio-temporal dropout pattern (MASK_RATIO), then force-unmasks any frame whose NDVI-vegetation pixel fraction is below CLOUD_GATE (a "don't waste masking on cloudy/non-vegetated frames" gate). We add a third, orthogonal masking stream over the spectral/band axis (SPECTRAL_MASK_P), forcing the encoder to learn inter-band correlation structure — the signal that separates visually-similar cereal classes (spring barley / winter triticale / mixed cereal), which dominate the baseline's mIoU deficit.
• NDVI auxiliary loss (NDVI_AUX_W) — a bounded phenology-index term added to the masked-reconstruction loss, so reconstructions have to be phenologically plausible, not just spectrally close in MSE.
• PrototypeHead (phenoproto.py) — replaces the linear classifier head with a cosine-distance classifier against per-class prototypes maintained as an EMA over a per-class memory bank (PROTO_BANK=4096, PROTO_MOMENTUM=0.999). Rare classes get a low-variance mean-based decision boundary instead of a linear weight that only ever sees frequency-proportional gradient.
• Balanced Supervised Contrastive loss (losses.py) — samples a fixed number of embedding anchors per present class (SUPCON_ANCHORS=64) rather than per pixel, so majority classes (e.g. Meadow) don't dominate the embedding geometry the way they dominate cross-entropy.
• Semi-supervised mean-teacher — an EMA teacher (EMA_DECAY=0.999) provides pseudo-labels on unlabeled fold-5 patches, gated by a FlexMatch-style per-class adaptive confidence threshold (AdaptiveThreshold, base PSEUDO_THRESH=0.95) so the pseudo-label distribution doesn't collapse onto the majority class.

Dataset

PASTIS — Sentinel-2 time series, 5 official geographic folds (with a 1km buffer between them to prevent spatial leakage). Splits follow the official STCLN protocol exactly (splits.py), including its two documented quirks: the labeled train/val sets are 76 hardcoded patch IDs each (not a fold filter), and pretrain/test use full folds.

Split	Source	Patches	Crops/epoch
Pretrain (Track A)	fold 5	496	7,936 (4×4 inner crop grid)
Train (finetune)	76 hardcoded IDs	76	152 (2 fixed crop positions)
Val (finetune)	76 hardcoded IDs	76	152
Test (Rung 0/A–D eval)	fold 4	482	7,712-equiv (full 128×128, no cropping)
Unlabeled (semi-sup)	fold 5	496	—

Verified programmatically at every launch (preflight.py): zero patch overlap between train/val/pretrain and the fold-4 test set.

Experimental protocol

• Normalization: per-fold Sentinel-2 band mean/std from PASTIS/NORM_S2_patch.json, averaged across folds.
• Positions: the model receives range(T) index positions, not real calendar dates (official behaviour — USE_INDEX_POSITIONS=True).
• Class weighting: CrossEntropy with background (class 0) and void (class 19) weighted to zero; both are also excluded from all reported metrics.
• Test-time protocol: fold-4 evaluation runs on the full 128×128 patch (not the 32×32 training crops) with no test-time augmentation, mIoU computed only over classes present in the ground truth — matching the official test_STCLN.py protocol exactly.

Configuration

Key hyperparameters (config.py); [OFFICIAL] = verbatim from STCLN, changing these invalidates the comparison to the published number, [NEW] = PhenoProto-SSL addition, [HW] = hardware tuning, no effect on optimization semantics.

Group	Param	Value
Model	N_CHANNELS / N_CLASSES	10 / 20	official
	ENCODER_WIDTHS / DECODER_WIDTHS	[32,256] / [32,256]	official
	AGG_MODE, N_HEAD, D_MODEL, D_K	att_mean, 8, 256, 32	official
Pretrain	PRE_EPOCHS / PRE_BATCH	100 / 4	official
	PRE_LR / PRE_WD / PRE_CLIP	1e-4 / 0.0 / 5.0	official
	MASK_RATIO / CLOUD_GATE / NDVI_THRESH	0.4 / 0.9 / 0.2	official
	SPECTRAL_MASK_P / NDVI_AUX_W	0.25 / 0.5	new
Finetune	FT_EPOCHS / FT_BATCH / FT_LR	100 / 2 / 1e-4	official
	USE_PROTOTYPE, D_EMBED, PROTO_TAU	True, 128, 0.1	new
	SUPCON_W, SUPCON_TAU, SUPCON_ANCHORS	0.1, 0.07, 64	new
	USE_SEMISUP, UNSUP_W, PSEUDO_THRESH	True, 1.0, 0.95	new
Reproducibility	SEED	3407	official
Hardware	AMP_DTYPE / NUM_WORKERS / TF32	bf16 / 12 / True	hw only

Full config with inline rationale for every value: config.py.

Hardware & environment

GPU	1× NVIDIA RTX PRO 6000 Blackwell Server Edition, 97,887 MiB VRAM
Driver / CUDA	580.159.03 / CUDA 13.0
PyTorch	2.13.0+cu130
Python	3.12.3
Precision	bf16 autocast (Blackwell/Ampere+: bf16 has fp32 exponent range, immune to the fp16 overflow that previously caused a silent NaN pretrain run)
Peak VRAM measured	~8.7 GB @ batch=4 (pretrain fwd/bwd, single 32×32 crop-tile step)

Status

As of the last check, run on this machine:

• Track A (S3M pretrain): ✅ complete, all 100 epochs, checkpoint_99 saved clean (0 non-finite weights).
• Track B — Rung 0 (official baseline): ✅ complete, all 100 epochs. Numbers below.
• Track B — Rung A_linear (official-checkpoint pipeline-equivalence check): reached epoch 34/100 (best val score 0.513) before being stopped for a methodology review; checkpoint preserved, not evaluated on the fold-4 test set. Rungs B_proto, C_supcon, D_full never started, and are intentionally not being resumed while the item below is open.
• s3m_ep89 — full PhenoProto-SSL finetune (prototype head + balanced SupCon + semi-supervised mean-teacher) sourced from our own Track A pretrain at epoch 89 (not the official released encoder): ✅ complete. Beats Rung 0 on the held-out test set — see Results.
• s3m_ep99 — same finetune config, sourced from the final (epoch 99) Track A pretrain checkpoint, kept as a separate run/tag from s3m_ep89: in progress.

An open methodology question is tracked before fully trusting any official-checkpoint-sourced Track B number (Rung 0, A_linear): the official finetuning_STCLN.py imports a src.dataset.PASTIS_Dataset from a sibling repo (utae-paps-main) that was never available in this environment and had to be reimplemented from the same logic already used elsewhere in this codebase (splits.py). Line-by-line fidelity against the true upstream implementation has not yet been confirmed. This does not affect s3m_ep89/s3m_ep99, which use our own PastisPatches loader throughout, not the vendored one.

Also worth knowing: s3m_ep89's validation mIoU peaked at epoch 13 (0.43) and steadily degraded to ~0.20 by epoch 99 (see the training-curve plot in Sample predictions) — the saved model_best.tar correctly captured the epoch-13 peak, so the reported result isn't affected, but training the full 100 epochs is currently counterproductive. Likely suspects: no LR decay (per official protocol) combined with the semi-sup loss weight ramping up over the same window. Not yet fixed.

Results

s3m_ep89 — full PhenoProto-SSL, our own S3M pretrain (epoch 89)

Fold-4 held-out test set, full-protocol evaluation, no TTA, best checkpoint (epoch 13/100 by validation mIoU):

Metric	s3m_ep89	Rung 0 (official baseline)	Published (STCLN_wp)
mIoU	0.4298	0.3493	0.4843
OA	0.7819	0.7746	0.8170
mF1	0.5488	0.4945	0.6059
Kappa	0.7328	0.7178	—

Beats the official-baseline reproduction on every metric, and unlike Rung 0 (4 dead classes), every one of the 18 scored classes gets a nonzero IoU. Still short of the published number, consistent with Rung 0 also falling short — see the open methodology question in Status.

Rung 0 — official STCLN reproduction (100/100 epochs)

Metric	This run	Published (STCLN_wp)	Δ
mIoU	0.3493	0.4843	−0.135
OA	0.7746	0.8170	−0.042
mF1	0.4945	0.6059	−0.111
Kappa	0.7178	—	—

Validation accuracy plateaued at epoch 7 (best checkpoint = epoch 7/100) and never improved for the remaining 92 epochs. 4 of 18 scored classes (Spring barley, Potatoes, Mixed cereal, Sorghum) are never predicted at all (IoU = 0) on the test set.

s3m_ep99 — full PhenoProto-SSL, our own S3M pretrain (final epoch 99)

Same finetune config as s3m_ep89, sourced from the final Track A checkpoint instead of the intermediate epoch-89 one. Fold-4 held-out test set, best checkpoint (epoch 14/100 by validation mIoU):

Metric	s3m_ep99	s3m_ep89	Rung 0 (official)	Published (STCLN_wp)
mIoU	0.4308	0.4298	0.3493	0.4843
OA	0.7873	0.7819	0.7746	0.8170
mF1	0.5525	0.5488	0.4945	0.6059
Kappa	0.7386	0.7328	0.7178	—

Our best result so far, essentially matching s3m_ep89 (marginal improvement) — consistent with pretraining another 10 epochs having a small but real effect. Same "zero dead classes" property as s3m_ep89. Its training curve (below) shows the same peak-then-degrade pattern as s3m_ep89 (peaks ~epoch 14 at mIoU 0.46, degrades afterward, with a sharp dip around epoch 62) — confirms this is a reproducible property of the current finetune recipe, not a one-off fluke.

Rungs A–D (PhenoProto-SSL ablations from the official checkpoint)

Not yet available — see Status.

Sample predictions

True-color composite / ground truth / prediction / error overlay on real PASTIS fold-4 test patches, s3m_ep89 model, generated by visualize_predictions.py (chosen as the 6 most class-diverse test patches — harder and more informative than random ones, most of which are >90% a single majority class):

Patch 40038 — mIoU 0.333, OA 0.723. Field boundaries and the two majority classes (Meadow/Corn) are largely correct; some confusion with Grapevine.

Patch 20368 — mIoU 0.164, OA 0.405. A genuinely hard patch (dense village + many thin/mixed parcels) — a real weakness, not cherry-picked.

More patches: patch_40019 · patch_40086 · patch_20156 · patch_20016

Training curve (validation mIoU/OA/mF1/Kappa per epoch) — shows the epoch-13 peak and subsequent degradation noted in Status:

s3m_ep99 (our best result), same 6 patches for direct comparison:

Patch 40038 — mIoU 0.428 (vs 0.333 for s3m_ep89), OA 0.796.

More s3m_ep99 patches: patch_40019 · patch_40086 · patch_20156 · patch_20016 · patch_20368

Regenerate for any checkpoint:

python3 visualize_predictions.py \
    --ckpt checkpoints/finetune/<tag>/model_best.tar --tag <tag> \
    --log logs/finetune_<tag>.log --n 6

Reproducing

# one-time setup: fixes PASTIS permissions, clones the official STCLN repo,
# checks Python deps and GPU/Blackwell compatibility
bash setup.sh

# mandatory gate — verifies splits are leak-free, class names match the
# official list, a real batch loads with correct shapes, pretrain/finetune
# forward+backward are finite, and reports peak VRAM. Exit 0 = safe to launch.
python3 preflight.py

# launch both tracks, detached in tmux
bash tmux_night.sh start

# any time, from any shell:
bash tmux_night.sh status      # GPU + both log tails + checkpoint progress
bash tmux_night.sh attach      # live view (Ctrl-b then d to detach)
bash tmux_night.sh results     # every metric found in the logs
bash tmux_night.sh stop        # kill the session

Every training script checkpoints every epoch (latest.tar) and auto-resumes from it on the next tmux_night.sh start — safe to interrupt for a VM pause/restart without losing more than one epoch of progress.

Environment overrides (for running on a different machine): PASTIS_ROOT, EXP_ROOT, REF_ROOT env vars override the dataset/exp/ reference-repo paths hardcoded as defaults in config.py.

Repository structure

phenoproto/
├── config.py               all paths + hyperparameters, inline rationale
├── splits.py                protocol-correct fold splits, class names, PastisPatches dataset
├── masking.py                S3M masking (official baseline + spectral stream) + reconstruction loss
├── phenoproto.py              PrototypeHead, PhenoProtoClassifier/Pretrain model wrappers
├── losses.py                 BalancedSupCon, Lovasz-softmax, AdaptiveThreshold, consistency loss
├── STCLN.py                  official UTAE/UTAEClassification backbone (vendored, unmodified)
├── pretrain_s3m.py            Track A entrypoint
├── finetune_phenoproto.py     Track B (Rungs A-D) entrypoint
├── evaluate_fold4.py          official fold-4 test protocol (--official flag for Rung 0's architecture)
├── visualize_predictions.py   true-color/GT/prediction/error maps + training-curve plots
├── preflight.py               pre-launch correctness/sanity gate
├── setup.sh                  one-time environment setup
├── track_b.sh                 Rung 0 + ablation ladder orchestrator
├── tmux_night.sh               overnight dual-track tmux orchestrator, resume-aware
└── logs/, checkpoints/, figures/  generated at runtime — see below

What to push to a public repo

This is a research project with a 37GB dataset and multi-GB checkpoints alongside the code — push code and documentation, not data or artifacts.

Push:

config.py  splits.py  masking.py  phenoproto.py  losses.py
pretrain_s3m.py  finetune_phenoproto.py  evaluate_fold4.py  preflight.py
visualize_predictions.py
setup.sh  track_b.sh  tmux_night.sh
README.md

Add before pushing (currently missing from this working copy):

• requirements.txt / environment.yml — pin torch, numpy, scikit-learn, tensorboard, pandas at minimum (see Hardware for the exact versions this was validated against).
• LICENSE — pick one; note that STCLN.py is vendored from XiaoleiQinn/STCLN and its own license terms should be checked/carried forward for that file specifically.
• .gitignore covering at minimum: __pycache__/, *.pyc, checkpoints/, logs/, cache/, *.tar.

Do not push:

• checkpoints/ — every .tar is 24–50MB; the pretrain dir alone is currently 683MB. Use Git LFS or a model registry (e.g. a separate HF Hub model repo) if checkpoints need to be distributed, not the code repo.
• logs/ — run logs and TensorBoard event files, fully regenerable.
• __pycache__/.
• The PASTIS dataset itself (37GB) — link to the official source instead.
• stcln_ref/ — this is a separate clone of the official STCLN repo (has its own .git) living alongside this project on disk, not a subdirectory of it; don't nest it into this repo. setup.sh clones it fresh.

Before pushing, also note: config.py's PASTIS_ROOT/EXP_ROOT/ REF_ROOT defaults are currently hardcoded to this machine's absolute paths (/home/ubuntu/IB-Connect-ver-2/test/...). They're already env-var-overridable (see Reproducing), but the fallback defaults leak this machine's path structure — harmless, but worth a find/replace pass to a placeholder before a public push.

Known issues & engineering notes

• Official repo required non-trivial repair to run at all. As cloned, finetuning_STCLN.py imported a src package from a sibling repo never included in the clone, hardcoded CUDA_VISIBLE_DEVICES='2' (breaks on any single-GPU box), hardcoded the dataset path inside main() regardless of its own --datadir flag, and needed tensorboard installed. It also crashed at runtime on pandas.DataFrame.append() (removed in pandas 3.0) and on a None-guard bug in its optional Visdom logger. All patched in-place with inline comments explaining each change; see git diff on stcln_ref/PASTIS/ for the full patch (not part of this repo — see What to push).
• CLOUD_GATE mostly neutralizes the spatio-temporal mask on real PASTIS data. The official masking gate force-unmasks any frame under 90% NDVI-vegetation coverage; measured on real pretrain-fold frames, the median vegetation fraction is 0.42, so the gate fires on ~98% of frames. CLOUD_GATE is left at its official value of 0.9 (unchanged) — this means S3M's effective tested contribution, as currently configured, is concentrated in the added spectral masking stream rather than the spatio-temporal stream. preflight.py's sanity bound on mask visible fraction was recalibrated from [0.35, 0.75] to [0.70, 1.0] to match this measured, expected behavior rather than an incorrect assumption.
• Batch size was tested and reverted. Raising PRE_BATCH (4→12→24) to use more of the 97GB card was tried; 24 nearly OOM'd with Track B running concurrently, and even the safe 12 measured slower wall-clock per epoch than 4 — this workload isn't GPU-compute-bound at this model/crop size. Reverted to the official value.

Acknowledgements

Backbone architecture and baseline training/eval protocol from XiaoleiQinn/STCLN. Dataset: PASTIS (Garnot & Landrieu).