OneScience-Group/SatMAE

<p align="center"> <strong> <span style="font-size: 30px;">SatMAE</span> </strong> </p>

Model Introduction

SatMAE is a masked autoencoder pre-training model for temporal and multispectral satellite imagery. It learns remote sensing image representations through temporal positional encoding, spectral group encoding, and independent masking across temporal or spectral dimensions, and is mainly used to improve performance on tasks such as satellite image classification, land-cover classification, multi-label classification, and semantic segmentation when labeled data is limited.

Paper: SatMAE: Pre-training Transformers for Temporal and Multi-Spectral Satellite Imagery https://arxiv.org/abs/2207.08051

Model Description

SatMAE was proposed by a research team at Stanford University. The model is pre-trained using fMoW satellite imagery, temporal satellite imagery, and Sentinel-2 multispectral imagery. The model is suitable for remote sensing tasks such as satellite image classification, land-cover classification, multi-label classification, and semantic segmentation.

Applicable Scenarios

Usage Instructions

1. OneCode Usage

Experience intelligent one-click AI4S programming through the OneCode online environment:

Experience intelligent one-click AI4S programming

2. Manual Installation and Usage

Hardware Requirements

• —GPU or DCU execution is recommended.
• —CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slow.
• —DCU users need to install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.

Download the Model Package

hf download OneScience-Group/SatMAE --local-dir ./SatMAE
cd SatMAE

Install the Runtime Environment

DCU Environment

# Activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# Supports installation with uv
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

GPU Environment

# Activate CONDA first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# Supports installation with uv
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

Training Data Introduction

The temporal experiment in the paper uses fMoW RGB, with input consisting of RGB image sequences of length 3 from the same location; the multispectral experiment uses fMoW-Sentinel. In temporal-mode NPZ files, images is [B,T,C,H,W], timestamps is [B,T,3], with the three fields being year - 2002, month - 1, and hour in order, and labels is [B]. The model also accepts continuous scalar time in [B,T] format.

Small synthetic data following the same protocol is used by default:

python scripts/fake_data.py

Run the command above when using synthetic data. When using real data, do not run fake_data.py; save the data as data/train.npz and data/test.npz, and modify conf/config.yaml according to the actual protocol.

images:     float32 [N,T,C,H,W]
timestamps: float32 [N,T,3]
labels:     int64   [N]

The three fields of timestamps are year - 2002, month - 1, and hour in order; continuous scalar time can also use the [N,T] format. The data should have completed size processing, channel ordering, temporal sorting, and numerical normalization.

Training

Single GPU:

python scripts/train.py

Multiple GPUs:

torchrun --nproc_per_node=8 scripts/train.py

Training outputs:

result/checkpoints/satmae.pt
result/training/metrics.json

The default configuration uses small synthetic data to validate the training workflow. Formal training should use the Base, Large, or Huge paper architecture presets provided by the code, together with real data and the paper's training budget.

Trained Weights

This repository provides weights trained on fMoW RGB temporal satellite imagery and fMoW-Sentinel multispectral satellite imagery in the weight/ folder. The weight files will be uploaded soon and are expected to be completed in the near future.

Inference

python scripts/inference.py

Inference results are output to:

result/output/reconstruction.npz

Evaluation and Visualization

python scripts/result.py

Evaluation and visualization outputs are saved to:

result/evaluation/metrics.json
result/evaluation/temporal_frame_reconstruction.png
result/evaluation/temporal_reconstruction_error.png
result/evaluation/spectral_band_reconstruction.png

Evaluation results include overall and masked-patch MSE, overall and masked-patch MSE for each temporal frame, and reconstruction MSE for each input channel. Synthetic-data results are only used to validate the engineering workflow and do not represent downstream transfer performance from the paper.

OneScience Official Information

Citation and License

This repository is a reproduction of the original SatMAE paper.