Swin-STCLN × PASTIS

Swin-STCLN: Hierarchical Swin Transformer Enhanced Spatio-Temporal Contrastive Learning Network for Crop Mapping

This repository implements an improved STCLN architecture for Sentinel-2 satellite image time-series semantic segmentation on the PASTIS benchmark.

The original STCLN pretrain → finetune workflow is preserved while replacing the flat CNN spatial encoder with a hierarchical Swin Transformer encoder and introducing cross-scale spatiotemporal fusion with boundary-aware refinement.

---

🏆 Results — 5-Fold Cross Validation

Summary

Official benchmark result:

mFscore = 58.35% ± 1.70%

Trained from scratch on AMD MI300X GPU.

---

📊 Per-Class IoU — All Folds

---

🏗️ Model Architecture — Swin-STCLN

1. Swin Spatial Encoder (SEncoder Replacement)

Original STCLN:

DoubleConv
+
DoubleConv

filters:
32 → 256

kernel:
3×3

Output:
(B×T,256,32,32)

Problem:

• flat representation
• single scale
• limited long range spatial modeling

Swin-STCLN replaces it:

Input:

(B,T,10,32,32)

Patch Embedding

Conv2D:

10 → 96 channels

kernel = 2

stride = 2

Spatial:

32×32 → 16×16

Output:

(B×T,96,16,16)

Swin Stage 1

2 Swin Transformer Blocks

• Window attention
• Shifted window attention

Configuration:

dim = 96

window = 4

Output:

f_fine:

(B×T,96,16,16)

Patch Merging

16×16 → 8×8

Channel:

96 → 192

Swin Stage 2

2 Swin Transformer Blocks

dim = 192

window = 4

Output:

f_coarse:

(B×T,192,8,8)

---

2. Temporal Encoder

The original STCLN Transformer Temporal Encoder is kept unchanged.

Configuration:

• Transformer layers: 3
• Attention heads: 8
• Sinusoidal positional encoding
• GroupNorm

Only input representation changes.

Original STCLN:

(B×32×32) × T × 256

Swin-STCLN:

(B×8×8) × T × 192

Benefits:

• reduces temporal tokens from 1024 → 64 locations
• temporal reasoning happens on semantic regions
• lower memory usage

---

3. STFusion — Cross Scale Fusion

Replacement for STCLN STA module.

Inputs:

Temporal branch:

coarse_agg

(B,192,8,8)

Spatial branch:

fine_agg

(B,96,16,16)

Process:

1. Upsample

8×8 →16×16

2. Projection

192 →96

3. Cross Attention

Query:

coarse semantic features

Key / Value:

fine spatial features

4. Residual fusion

5. Upsample

16×16 →32×32

Output:

(B,128,32,32)

---

4. Pretraining Reconstruction Decoder

The original STCLN self-supervised learning objective is preserved.

Unchanged:

• Spatiotemporal masking
• Mask ratio = 0.4
• MSE reconstruction loss
• Unlabeled pretraining patches = 7936

Difference:

Base STCLN reconstructs directly using a Linear layer because temporal features remain at 32×32.

Swin-STCLN reconstructs from hierarchical features.

TEncoder output:

(B,T,192,8,8)

Reconstruction Head:

ConvTranspose2D

8×8
 ↓
16×16


ConvTranspose2D

16×16
 ↓
32×32


Conv2D

192 → 10 Sentinel-2 bands

Final reconstruction:

(B,T,10,32,32)

The same masked pixel MSE objective is applied.

---

5. Finetuning Decoder

The original STCLN linear decoder is replaced with a boundary-aware segmentation decoder.

Input:

STFusion output

(B,128,32,32)

Semantic Decoder

Architecture:

Conv-BN-ReLU

128 → 64

Conv-BN-ReLU

64 →64

1×1 Conv classifier

Output:

(B,18,32,32)

---

Boundary Decoder

Parallel boundary prediction branch:

Conv-BN-ReLU

128 →64

Conv-BN-ReLU

64 →64

1×1 Conv

Output:

(B,1,32,32)

Boundary supervision is generated automatically using morphological gradient from semantic masks.

No additional annotation required.

---

Gated Refinement

Semantic feature

Boundary feature

↓

Concatenation

↓

Convolution refinement

↓

Final refined prediction:

(B,18,32,32)

---

6. Finetuning Loss

Original STCLN:

Cross Entropy

Swin-STCLN:

Total Loss =

CE(semantic output)

+

CE(refined output)

+

0.5 × BCE(boundary output)

The boundary loss improves separation between neighbouring crop parcels.

Boundary weight = 0.5

---

⚡ Inference Speed

Measured on AMD MI300X GPU.

---

📁 Repository Structure

Swin-STCLN-PASTIS/

├── models/

│   ├── swin_encoder.py
│       # PatchEmbed + Swin Blocks + PatchMerging

│   ├── temporal_encoder.py
│       # STCLN Transformer Encoder

│   ├── stfusion.py
│       # Cross-scale attention fusion

│   ├── decoder.py
│       # Semantic decoder
│       # Boundary decoder
│       # Gated refinement

│   ├── reconstruction.py
│       # ConvTranspose reconstruction head

│   └── swin_stcln.py
│       # Complete architecture


├── datasets/

│   └── pastis_dataset.py


├── losses/

│   ├── segmentation_loss.py
│   └── boundary_loss.py


├── evaluation/

│   └── metrics.py


├── train.py

├── pretrain.py

├── finetune.py

├── visualize_results.py

├── checkpoints/

└── results/

---

🚀 Quick Start

Installation

git clone https://huggingface.co/Dhruv1000/Swin-STCLN-PASTIS

cd Swin-STCLN-PASTIS


pip install torch torchvision timm einops geopandas matplotlib scikit-learn

---

Training

Single fold:

python train.py \
    --data_root /path/to/PASTIS \
    --fold 1 \
    --epochs 100 \
    --batch_size 16 \
    --lr 5e-5 \
    --weight_decay 0.05 \
    --warmup_iters 500 \
    --num_workers 4 \
    --amp \
    --work_dir ./work_dirs/fold1

---

5-Fold Cross Validation

for fold in 1 2 3 4 5
do

python train.py \
    --data_root /path/to/PASTIS \
    --fold $fold \
    --epochs 100 \
    --batch_size 16 \
    --lr 5e-5 \
    --work_dir ./work_dirs/fold${fold}

done

---

Inference

import torch

from models.swin_stcln import build_swin_stcln


model = build_swin_stcln(
    num_classes=18
)


checkpoint = torch.load(
    "checkpoints/best_model.pth",
    weights_only=False
)


model.load_state_dict(
    checkpoint["model"]
)


model.eval()



# Input:
# Batch
# Time
# Sentinel-2 Bands
# Height
# Width


x = torch.randn(
    1,
    32,
    10,
    32,
    32
)


logits = model(x)


# Output:

# (1,18,32,32)


prediction = logits.argmax(dim=1)

---

📋 Training Configuration

---

📦 Dataset — PASTIS

The model is evaluated on the PASTIS (Panoptic Agricultural Satellite Time Series) benchmark.

The official train / validation / test split is preserved.

---

🌾 PASTIS Classes

---

🔥 What Remains Identical To Original STCLN

The following components are unchanged:

✅ Pretrain → finetune workflow

✅ Spatiotemporal masked reconstruction

✅ Mask ratio:

0.4

✅ Reconstruction objective:

Mean Squared Error

✅ Temporal Encoder design:

• 3 Transformer layers
• 8 attention heads
• sinusoidal positional encoding
• GroupNorm

✅ Dataset protocol:

• PASTIS32
• Sentinel-2
• 10 spectral channels
• Official folds

✅ Weight transfer:

Pretrained:

SEncoder
+
TEncoder


        ↓


Finetuning initialization

---

🆚 STCLN vs Swin-STCLN

---

📈 Training Dynamics (Fold 1)

Best checkpoint:

Epoch 92

Total training time per fold:

~32 minutes on AMD MI300X

---

🖼️ Visualizations

Generated evaluation plots:

Per Fold

results/fold{N}/plots/

Includes:

• Training curves
• Per-class IoU plots
• Metrics radar chart
• Confusion matrix
• Prediction maps
• Error maps
• IoU scatter analysis
• Overfitting analysis

---

🔧 Implementation Details

Pure PyTorch Implementation

No dependency on:

• MMSegmentation
• MMEngine
• external training framework

Implemented with:

• PyTorch
• timm
• einops

---

Major Architectural Changes

1. Hierarchical Spatial Learning

CNN encoder replaced by Swin Transformer blocks.

Benefits:

• larger receptive field
• window attention
• shifted window information exchange

---

2. Efficient Temporal Modelling

Instead of:

1024 temporal sequences/image

Swin-STCLN uses:

64 temporal sequences/image

This reduces memory while keeping semantic information.

---

3. Cross Scale Feature Recovery

The coarse temporal representation loses boundaries.

STFusion restores:

• high level semantics
• fine spatial details

using cross attention fusion.

---

4. Boundary Aware Learning

Boundary decoder learns crop separation using automatically generated masks.

No manual boundary annotation needed.

---

📚 Citation

If this implementation is useful, cite the original STCLN work and PASTIS benchmark.

@inproceedings{garnot2021pastis,

title={Panoptic Segmentation of Satellite Image Time Series with Convolutional Temporal Attention Networks},

author={Garnot, Vivien Sainte Fare and Landrieu, Loic},

booktitle={ICCV},

year={2021}

}

For Swin Transformer:

@inproceedings{liu2021swin,

title={Swin Transformer: Hierarchical Vision Transformer using Shifted Windows},

author={Liu, Ze and Lin, Yutong and Cao, Yue and others},

booktitle={ICCV},

year={2021}

}

---

📄 License

Apache-2.0

---

Trained on AMD MI300X
ROCm 7.0
PyTorch 2.x

Swin-STCLN × PASTIS

Hierarchical Swin Spatial Encoder
+ STCLN Temporal Encoder
+ Cross Scale Boundary-Aware Fusion