README.md

---
license: apache-2.0

language:
- en

tags:
- remote-sensing
- semantic-segmentation
- agriculture
- crop-mapping
- sentinel-2
- video-swin-transformer
- pastis
- pytorch
- amd-rocm
- foundation-model

datasets:
- pastis

metrics:
- mean_iou
- f1
- accuracy

model-index:
- name: AgriFM-small-PASTIS
  results:
  - task:
      type: image-segmentation
      name: Semantic Segmentation
    dataset:
      name: PASTIS
      type: pastis
    metrics:
    - type: mean_iou
      value: 44.43
      name: Mean IoU (5-fold CV)
    - type: f1
      value: 58.35
      name: F1 Score (5-fold CV)
    - type: accuracy
      value: 68.36
      name: Overall Accuracy (5-fold CV)
---

# AgriFM × PASTIS

> **Reimplementation of AgriFM** — *A Multi-source Temporal Remote Sensing Foundation Model for Agriculture Mapping* (Li et al., 2025) — adapted for the [PASTIS](https://github.com/VSainteuf/pastis-benchmark) benchmark dataset and trained from scratch on a single AMD MI300X GPU (ROCm 7.0).

---

## 🏆 Results — 5-Fold Cross-Validation

### Summary

| Metric | Fold 1 | Fold 2 | Fold 3 | Fold 4 | Fold 5 | **Mean ± Std** |
|--------|--------|--------|--------|--------|--------|----------------|
| **mFscore** | 58.37% | 56.69% | 59.18% | 61.00% | 56.49% | **58.35% ± 1.70%** |
| **mIoU** | 44.35% | 42.90% | 45.54% | 46.80% | 42.58% | **44.43% ± 1.62%** |
| **OA** | 67.10% | 66.85% | 69.96% | 71.48% | 66.41% | **68.36% ± 2.13%** |
| **Kappa** | 60.02% | 59.71% | 62.79% | 64.33% | 58.99% | **61.17% ± 2.19%** |
| **mPrecision** | 52.78% | 51.33% | 54.51% | 57.40% | 51.24% | **53.45% ± 2.47%** |
| **mRecall** | 70.72% | 69.00% | 68.69% | 68.93% | 68.72% | **69.21% ± 0.82%** |

> **Official benchmark result: mFscore = 58.35% ± 1.70%** — trained from scratch, no pretrained weights, 32 minutes per fold on AMD MI300X.

---

## 📊 Per-Class IoU — All Folds

| Class | Fold 1 | Fold 2 | Fold 3 | Fold 4 | Fold 5 | **Mean** |
|-------|--------|--------|--------|--------|--------|----------|
| 🌽 Corn | 76.76% | 75.31% | 75.75% | 77.09% | 74.60% | **75.90%** |
| 🌿 Winter rapeseed | 74.17% | 76.05% | 77.14% | 77.49% | 72.90% | **75.55%** |
| 🌾 Beet | 73.32% | 72.52% | 78.01% | 71.87% | 73.78% | **73.90%** |
| 🌾 Soft winter wheat | 72.22% | 68.92% | 71.64% | 71.18% | 73.37% | **71.47%** |
| 🌿 Soybeans | 66.81% | 55.03% | 63.35% | 62.73% | 59.49% | **61.48%** |
| 🌾 Winter barley | 56.56% | 59.08% | 57.70% | 62.90% | 56.94% | **58.64%** |
| 🌱 Meadow | 56.56% | 55.96% | 57.47% | 58.18% | 55.29% | **56.69%** |
| 🌻 Sunflower | 53.94% | 48.13% | 64.54% | 57.16% | 50.59% | **54.87%** |
| 🟫 Background | 48.39% | 50.20% | 51.60% | 54.88% | 49.83% | **50.98%** |
| 🌾 Winter durum wheat | 43.30% | 43.68% | 49.19% | 47.26% | 39.38% | **44.56%** |
| 🥔 Potatoes | 37.33% | 37.75% | 34.93% | 42.60% | 29.07% | **36.34%** |
| 🌿 Grapevine | 36.16% | 36.62% | 37.52% | 38.60% | 34.60% | **36.70%** |
| 🌾 Spring barley | 30.51% | 40.00% | 32.87% | 38.06% | 33.97% | **35.08%** |
| 🌿 Leguminous fodder | 23.17% | 20.88% | 21.02% | 28.88% | 18.12% | **22.41%** |
| 🍎 Fruits/veg/flowers | 22.27% | 19.75% | 20.97% | 22.67% | 21.97% | **21.53%** |
| 🌾 Winter triticale | 25.01% | 19.33% | 24.30% | 23.72% | 19.32% | **22.34%** |
| 🌾 Mixed cereal | 17.63% | 12.64% | 13.20% | 16.81% | 13.02% | **14.66%** |
| 🍑 Orchard | 15.25% | 13.11% | 26.05% | 18.81% | 13.14% | **17.27%** |
| 🌿 Sorghum | 13.29% | 10.12% | 8.04% | 18.40% | 19.68% | **13.91%** |

---

## ⚡ Inference Speed

Measured on AMD MI300X (191.7 GB VRAM), AgriFM-small (39.6M params):

| Batch Size | Time (ms) | Throughput | VRAM Used |
|-----------|-----------|------------|-----------|
| 1 | 7.8 ms | 128.9 patches/sec | 1.12 GB |
| 4 | 19.6 ms | 203.8 patches/sec | 2.08 GB |
| 8 | 37.1 ms | 215.4 patches/sec | 3.40 GB |
| 16 | 73.4 ms | 218.0 patches/sec | 6.03 GB |
| 32 | 143.0 ms | 223.8 patches/sec | 11.31 GB |
| 64 | 281.0 ms | 227.8 patches/sec | 21.85 GB |

> Each patch = 32 temporal frames × 10 spectral bands × 128×128 pixels of Sentinel-2 imagery.

---

## 🏗️ Model Architecture

```
Input: (B, 32, 10, 128, 128)  ← 32 frames, 10 S2 bands, 128×128 pixels
    ↓
SwinPatchEmbed3D  patch_size=(4,2,2)  → (B, 64, 8, 64, 64)
    ↓
SwinTransformer3D — 4 stages with synchronized spatiotemporal downsampling
  Stage 1: depth=2, heads=2,  dim=64   → (B, 64,  8, 64, 64)
  Stage 2: depth=2, heads=4,  dim=128  → (B, 128, 4, 32, 32)
  Stage 3: depth=6, heads=8,  dim=256  → (B, 256, 2, 16, 16)
  Stage 4: depth=2, heads=16, dim=512  → (B, 512, 1,  8,  8)
    ↓
feature_fusion='cat' → (B, 512, 8, 8)  ← temporal dim collapsed into channels
    ↓
MultiFusionNeck  — 3-level U-Net decoder with skip connections
  Level 1: 8×8   → 16×16  (skip from Stage 3)
  Level 2: 16×16 → 32×32  (skip from Stage 2)
  Level 3: 32×32 → 64×64  (skip from Stage 1)
  Final:   64×64 → 128×128
    ↓
CropFCNHead  Conv(512→256)+ReLU+Conv(256→20)
    ↓
Output: (B, 20, 128, 128)  ← per-pixel crop type probabilities
```

### Key Innovation: Synchronized Spatiotemporal PatchMerging

Unlike standard Video Swin Transformer which only downsamples spatially, AgriFM's `PatchMerging` with `mean_frame_down=True` simultaneously:
- **Spatially**: 4 strided samples concatenated → 4C channels
- **Temporally**: mean pooling over D_step consecutive frames

This keeps spatial and temporal scales synchronized throughout the hierarchy.

---

## 📁 Repository Structure

```
AgriFM_PASTIS/
├── models/
│   ├── video_swin_transformer.py   # VST backbone + synchronized PatchMerging
│   ├── neck.py                     # MultiFusionNeck (U-Net decoder)
│   ├── heads.py                    # CropFCNHead
│   └── agrifm.py                   # Full model + build_agrifm_pastis_small/tiny/base
├── datasets/
│   └── pastis_dataset.py           # Official 5-fold splits via geopandas
├── losses/
│   └── loss.py                     # CropCELoss + hard mining + class weights
├── evaluation/
│   └── metrics.py                  # OA, mIoU, mFscore, Precision, Recall, Kappa
├── train.py                        # Training loop with colored output + logging
├── visualize_results.py            # 8-plot visualization suite per fold
├── visualize_all_folds.py          # Cross-validation summary plots
├── requirements.txt
├── checkpoints/
│   └── fold1_v3_best_model.pth     # Best checkpoint: fold1, epoch 92, mFscore=58.37%
└── results/
    ├── fold1/                      # Logs, metrics, plots for fold 1
    ├── fold2/ ... fold5/           # Same for folds 2–5
    └── all_folds_summary/          # Cross-validation summary plots
```

---

## 🚀 Quick Start

### Installation

```bash
git clone https://huggingface.co/Dhruv1000/AgriFM_PASTIS
cd AgriFM_PASTIS
pip install torch torchvision timm einops geopandas matplotlib seaborn scikit-learn
```

### Training (single fold)

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

### Training (all 5 folds)

```bash
for fold in 1 2 3 4 5; do
    python train.py \
        --data_root /path/to/PASTIS \
        --fold $fold \
        --small_model \
        --epochs 100 \
        --batch_size 16 \
        --lr 5e-5 \
        --weight_decay 0.05 \
        --work_dir ./work_dirs/fold${fold}
done
```

### Inference

```python
import torch
from models.agrifm import build_agrifm_pastis_small

model = build_agrifm_pastis_small(num_classes=20)
ckpt  = torch.load('checkpoints/fold1_v3_best_model.pth',
                   weights_only=False)
model.load_state_dict(ckpt['model'])
model.eval()

# Input: (B, T, C, H, W) = (batch, 32 frames, 10 bands, 128px, 128px)
x      = torch.randn(1, 32, 10, 128, 128)
logits = model(x)   # (1, 20, 128, 128)
pred   = logits.argmax(dim=1)  # (1, 128, 128) — class indices
```

### Visualize Results

```bash
# Per-fold visualization (8 plots)
python visualize_results.py \
    --work_dir ./work_dirs/fold1 \
    --data_root /path/to/PASTIS \
    --fold 1 --model_size small

# All-folds cross-validation summary
python visualize_all_folds.py
```

---

## 📋 Training Configuration

| Parameter | Value |
|-----------|-------|
| Model | AgriFM-small (39.6M params) |
| Optimizer | AdamW (β=(0.9, 0.999)) |
| Learning rate | 5e-5 |
| Weight decay | 0.05 |
| LR schedule | Linear warmup (500 iters) → cosine decay to 1e-6 |
| Batch size | 16 |
| Epochs | 100 |
| Loss | Cross-entropy + hard mining (top 50%) + inverse-freq class weights |
| Augmentation | Random h-flip, v-flip, 90/180/270° rotation |
| AMP | Enabled (torch.amp, ROCm 7.0 compatible) |
| Gradient clipping | max_norm = 5.0 |
| Input frames | 32 (uniform sampled from 61) |
| Normalization | Per-fold mean/std from NORM_S2_patch.json |

---

## 📦 Dataset: PASTIS

| Property | Details |
|----------|---------|
| Total patches | 2,433 geo-referenced tiles |
| Patch size | 128 × 128 pixels at 10m resolution |
| Satellite | Sentinel-2 (10 spectral bands) |
| Temporal length | 61 observations per patch |
| Label hierarchy | 3 levels — we use channel 0 (20 coarse classes) |
| Folds | 5 official geographic folds (~487–496 patches each) |
| Fold assignment | From `metadata.geojson` via geopandas |
| Ignore index | Class 19 (Void) excluded from all metrics |

### PASTIS Classes (20 coarse classes)

| ID | Class | Avg IoU (5-fold) |
|----|-------|-----------------|
| 0 | Background | 50.98% |
| 1 | Meadow | 56.69% |
| 2 | Soft winter wheat | 71.47% |
| 3 | Corn | 75.90% |
| 4 | Winter barley | 58.64% |
| 5 | Winter rapeseed | 75.55% |
| 6 | Spring barley | 35.08% |
| 7 | Sunflower | 54.87% |
| 8 | Grapevine | 36.70% |
| 9 | Beet | 73.90% |
| 10 | Winter triticale | 22.34% |
| 11 | Winter durum wheat | 44.56% |
| 12 | Fruits/veg/flowers | 21.53% |
| 13 | Potatoes | 36.34% |
| 14 | Leguminous fodder | 22.41% |
| 15 | Soybeans | 61.48% |
| 16 | Orchard | 17.27% |
| 17 | Mixed cereal | 14.66% |
| 18 | Sorghum | 13.91% |
| 19 | Void (ignore) | — |

---

## 🔧 Key Implementation Decisions

### 1. No MMSeg/MMEngine dependency
Original GitHub uses MMSegmentation. This reimplementation uses **pure PyTorch only** — no registry system, no config framework. Runs anywhere PyTorch runs.

### 2. Official fold splits
Initial fold assignment by patch ID prefix caused catastrophic distribution mismatch (W.Durum ratio: 88× between train/val). Fixed by loading official fold assignments from `metadata.geojson` using geopandas. Val loss dropped from 34 → 0.67 at epoch 1.

### 3. Class-weighted loss
Added `compute_class_weights()` with inverse-frequency weighting. Background weight = 0.01, Winter Durum Wheat weight = 7.24 — essential for learning rare crops.

### 4. Model size variants

| Variant | Params | embed_dim | depths | Use case |
|---------|--------|-----------|--------|----------|
| `base` | 196.2M | 128 | [2,2,18,2] | Paper-faithful, needs pretrained weights |
| `small` ✅ | 39.6M | 64 | [2,2,6,2] | **Used here — recommended for scratch** |
| `tiny` | 9.6M | 32 | [2,2,4,2] | Backup if small overfits |

### 5. AMD ROCm compatibility
- `torch.amp.GradScaler('cuda')` and `torch.amp.autocast('cuda')` (new API)
- `torch.load(..., weights_only=False)` for numpy scalar checkpoints
- `MIOPEN_LOG_LEVEL=0` to suppress kernel selection warnings (not errors)

### 6. Added Cohen's Kappa metric
Not in original codebase. Computed from full confusion matrix. Final 5-fold mean Kappa = **61.17%** (substantial agreement).

---

## 🆚 Comparison: Paper vs GitHub vs This Implementation

| Aspect | Paper | GitHub | **Ours** |
|--------|-------|--------|----------|
| Framework | PyTorch | PyTorch + MMSeg + MMEngine | **Pure PyTorch** |
| Dataset | 25.2M global images | Custom HDF5 | **PASTIS .npy** |
| Input size | 224×224 | 256×256 | **128×128** |
| Patch embed | SwinPatchEmbed3D | Identical | **Identical ✓** |
| Patch size | (4,2,2) | (4,2,2) | **Identical ✓** |
| PatchMerging | Synchronized | mean_frame_down=True | **Identical ✓** |
| feature_fusion | Not explicit | 'cat' | **Identical ✓** |
| Window size | Not explicit | (8,7,7) | **Identical ✓** |
| Neck | MultiFusionNeck | MultiFusionNeck | **Identical ✓** |
| Head | CropFCNHead | CropFCNHead | **Identical ✓** |
| Pre-training | LC fractions + L1 | Same | **Not implemented** |
| Class weights | Not mentioned | Not implemented | **Added ✓** |
| Kappa metric | Not reported | Not implemented | **Added ✓** |
| Augmentation | Not detailed | flip/rotate | **flip + rotate ✓** |
| Grad clipping | Not specified | None | **max_norm=5.0 ✓** |
| GPU | 10× NVIDIA L40 | 10× NVIDIA L40 | **1× AMD MI300X** |

---

## 📈 Training Dynamics (Fold 1 example)

| Epoch | Train Loss | Val Loss | mFscore | mIoU | Kappa |
|-------|-----------|----------|---------|------|-------|
| 1 | 0.878 | 0.671 | 2.79% | 1.47% | 2.39% |
| 4 | 0.431 | 0.472 | 20.08% | 12.34% | 10.76% |
| 10 | 0.320 | 0.380 | ~33% | ~22% | ~24% |
| 18 | 0.222 | 0.323 | 35.55% | 24.17% | 24.93% |
| 55 | 0.083 | 0.363 | 53.37% | 39.92% | 53.16% |
| **92** | **0.050** | **0.350** | **58.20%** | **44.20%** | **60.0%** |
| 100 | 0.048 | 0.360 | 57.90% | 44.10% | 59.8% |

Total training time per fold: **~32 minutes** on AMD MI300X.

---

## 🖼️ Visualizations

The repository includes 8 visualization plots per fold plus 6 cross-validation summary plots:

**Per-fold (in `results/fold{N}/plots/`):**
- `0_summary_card.png` — Dark-theme dashboard with all metrics
- `1_training_curves.png` — 6-panel training history
- `2_per_class_iou.png` — Sorted horizontal bar chart
- `3_metrics_radar.png` — Spider/radar chart of all 6 metrics
- `4_confusion_matrix.png` — Raw + normalized confusion matrices
- `5_prediction_maps.png` — 6 sample RGB / GT / Pred / Error maps
- `6_class_iou_scatter.png` — Per-class IoU scatter with mIoU line
- `7_overfitting_analysis.png` — Train vs val loss gap over epochs

**Cross-validation (`results/all_folds_summary/`):**
- `CV1_metrics_all_folds.png` — Grouped bar chart all folds
- `CV2_per_class_heatmap.png` — IoU heatmap: folds × classes
- `CV3_training_curves_overlay.png` — All folds on same axes
- `CV4_summary_table.png` — Mean ± std table
- `CV5_mean_per_class_iou.png` — Mean IoU with std error bars
- `CV6_metric_boxplots.png` — Box plots with individual fold points

---

## 📚 Citation

If you use this code or results, please cite the original AgriFM paper:

```bibtex
@article{li2025agrifm,
  title     = {AgriFM: A Multi-source Temporal Remote Sensing Foundation Model
               for Agriculture Mapping},
  author    = {Li, Wenyuan and Liang, Shunlin and Chen, Kun and others},
  journal   = {Remote Sensing of Environment},
  year      = {2025},
  doi       = {10.1016/j.rse.2026.115234}
}
```

And the PASTIS dataset:

```bibtex
@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}
}
```

---

## 📄 License

Apache 2.0 — see LICENSE for details.

---

*Trained on AMD MI300X · ROCm 7.0 · PyTorch 2.x · April 2026*