JEDI: JEPA-to-Edge Distillation for Efficient Cropland Segmentation from Satellite Imagery

Official code for JEDI, a two-stage cross-architecture knowledge distillation framework that transfers an I-JEPA (ViT-H/16) representation into a compact SegFormer student for cropland segmentation from Sentinel-2 imagery.

---

Overview

Large self-supervised vision models provide strong representations for remote-sensing segmentation but are too expensive to deploy at the satellite or field edge. JEDI distills such a representation from a large isotropic I-JEPA ViT teacher into a compact hierarchical SegFormer student in two stages:

• Stage 1 — Representation Alignment. The student encoder is trained, with no segmentation labels, to reproduce the frozen teacher's latent token representation through a lightweight cross-architecture alignment module (linear projection + spatial resampling).
• Stage 2 — Tri-Objective Joint Refinement. The aligned student is jointly optimized with (i) supervised segmentation loss, (ii) temperature-scaled response distillation from the teacher's predictions, and (iii) a persistent feature-alignment term that remains active throughout task adaptation, rather than being dropped after pretraining as in conventional pipelines.

At inference, only the student encoder and decode head are deployed; the teacher, its decoder, and the alignment module are discarded.

---

Repository Contents

This repository contains the training and evaluation code of JEDI.

JEDI/
├── train.py          # two-stage training script (Stage 1 + Stage 2)
├── assets/
│   └── architecture.png
├── LICENSE
└── README.md

Figure-generation, checkpoint-inference, and benchmarking utilities used to produce the paper's figures and tables are not included in this repository.

No pretrained checkpoints or preprocessed data are distributed with this repository. To reproduce results, you will need to (1) obtain the CalCROP21 dataset and a fine-tuned I-JEPA teacher checkpoint yourself, and (2) run training from scratch using train.py below.

---

Requirements

This code depends on standard deep learning and remote-sensing libraries, including:

• Python 3.9+
• PyTorch
• Transformers (IJepaModel, SegformerForSemanticSegmentation)
• torchmetrics
• NumPy
• tqdm

---

Data Preparation

train.py expects the CalCROP21 directory layout:

<root_path>/
├── IMAGE_GRIDS/
│   └── <tile_name>_IMAGE.npy      # (T, C, H, W) Sentinel-2 stack
├── STATT_GRIDS/
│   └── <tile_name>_STATT.npy      # (H, W) CDL-derived crop labels
└── GRID_SPLITS/
    ├── TRAIN_SET.npy              # array of tile-name strings
    ├── TEST_SET.npy
    └── VALIDATION_SET.npy

Cropland labels are binarized as label < 22 (cultivated vs. non-cultivated). Each multi-temporal tile is reduced to a single three-channel composite by selecting the acquisition with maximal mean NDVI before training and inference.

You will also need a fine-tuned I-JEPA teacher checkpoint (--teacher_ckpt), containing a state_dict with model.-prefixed keys, as produced by fine-tuning facebook/ijepa_vith16_1k with a convolutional decoder head for binary segmentation.

---

Usage

Run the full two-stage pipeline (Stage 1 runs automatically if no Stage 1 checkpoint is found; otherwise it is skipped):

python3 train.py \
  --root_path /path/to/ACCEPTABLE_GRIDS \
  --splits_folder /path/to/ACCEPTABLE_GRIDS/GRID_SPLITS \
  --teacher_ckpt /path/to/teacher_checkpoint.ckpt \
  --segformer_variant b0 \
  --save_dir models \
  --device_id 0

---

License

This project is released under the MIT License.