OneScience-Group/SatMAE-PP

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

Model Introduction

SatMAE++ is a masked autoencoder for optical and multispectral satellite imagery. It uses a visible-token encoder to learn remote sensing representations and reconstructs native spatial-scale targets through a convolutional multiscale decoder. The model supports RGB and grouped Sentinel inputs, with a training objective composed of MSE+L1 for masked patches and MSE+L1 for multiscale reconstruction.

Paper: Rethinking Transformers Pre-training for Multi-Spectral Satellite Imagery https://arxiv.org/abs/2403.05419

Model Description

SatMAE++ was proposed by the team of Mubashir Noman, Muzammal Naseer, Hisham Cholakkal, Rao Muhammad Anwar, Salman Khan, and Fahad Shahbaz Khan at Mohamed bin Zayed University of Artificial Intelligence in Abu Dhabi. The model is trained using FMoW-RGB and FMoW-Sentinel imagery. The model is suitable for representation learning and multiscale reconstruction tasks involving optical and multispectral satellite imagery.

Applicable Scenarios

Usage Instructions

1. OneCode Usage

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

Experience intelligent AI4S programming

2. Manual Installation and Usage

Hardware Requirements

• —GPU or DCU execution is recommended.
• —CPU can be used for small-configuration connectivity validation.
• —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-PP --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 paper uses FMoW-RGB and FMoW-Sentinel for pre-training. In NPZ files, images is [B,C,H,W], with native images_2x or images_4x targets consistent with the configured scales. RGB uses three channels, while Sentinel uses ten channels encoded by spectral groups.

Synthetic data 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; organize the data as data/train.npz and data/test.npz, where RGB data contains at least:

images:    float32 [N,3,H,W]
images_2x: float32 [N,3,2H,2W]

Sentinel data uses ten channels and provides images_2x and images_4x according to the configuration:

images:    float32 [N,10,H,W]
images_2x: float32 [N,10,2H,2W]
images_4x: float32 [N,10,4H,4W]

The high-resolution fields should be native targets registered with the input scene and should not be replaced by temporary interpolation of low-resolution imagery. After preparing the data, modify conf/config.yaml according to the actual number of channels, input size, spectral grouping, and training scales.

Training

python scripts/train.py

Multiple GPUs:

torchrun --nproc_per_node=8 scripts/train.py

Training outputs:

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

Training outputs include a model checkpoint that can be used for subsequent inference and metric records reflecting the training process and loss convergence, facilitating training-state preservation and analysis of model optimization.

The default configuration is a quick synthetic-data protocol. The formal paper configuration for RGB is ViT-L, 224/16, and 800 epochs; the Sentinel configuration is ViT-L, 96/8, 10-channel three-group input after removing B1/B9/B10, and 50 epochs.

Trained Weights

This repository provides weights trained on FMoW-RGB optical 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/multiscale_reconstruction.png
result/evaluation/scale_comparison.png

Evaluation results measure model performance in terms of overall reconstruction error, masked-region recovery quality, and reconstruction error at different native scales, and intuitively show the model's multiscale recovery capability through multiscale reconstruction results and cross-scale error comparisons.

OneScience Official Information

Citation and License

This repository is a reproduction of the original SatMAE++ paper.