OneScience-Group/SatlasPretrain

SatlasPretrain

Model Introduction

SatlasPretrain is a large-scale remote sensing pre-training approach for learning from multi-temporal imagery and seven label types. Its unified SatlasNet model supports semantic segmentation, regression, point detection, polygon instances, polylines, object properties, and image classification for Earth observation and downstream remote sensing applications.

Paper: SatlasPretrain: A Large-Scale Dataset for Remote Sensing Image Understanding

https://arxiv.org/abs/2211.15660

Model Description

SatlasPretrain was proposed by the Allen Institute for AI and trained on Sentinel-2 and NAIP imagery with 302 million labels in 137 categories and seven label types. The model is suitable for multi-temporal remote sensing segmentation, regression, detection, property prediction, and classification tasks.

Use Cases

Scenario	Description
Remote sensing multi-task training	Train SatlasPretrain with multi-temporal imagery and segmentation, regression, detection, property, and classification labels.
Local quick validation	Use synthetic data to validate loading, multi-task training, inference, evaluation, and qualitative visualization.
Hugging Face / OneCode execution	Download the standalone model package, install dependencies, and run the scripts directly.
Multi-GPU training	Launch distributed multi-process training with torchrun.

Usage Guide

1. OneCode Usage

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

Click to Experience Intelligent One-Click AI4S Programming

2. Manual Installation and Usage

Hardware Requirements

• A GPU or DCU is recommended.
• CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
• DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching the cluster, is recommended.

Download the Model Package

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

Install the Runtime Environment

DCU Environment

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

GPU Environment

# Please 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
# uv installation is supported
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

Training Data Introduction

The original SatlasPretrain dataset is published by the Allen Institute for AI and organized as Web-Mercator tiles. Sentinel-2, Sentinel-1, and NAIP imagery and bands are stored as PNG files. Point, polygon, polyline, property, and classification labels are stored in vector.json; semantic segmentation and regression labels are grayscale PNG files; split and image-time metadata are JSON files.

Official data resources:

• https://github.com/allenai/satlas/blob/main/SatlasPretrain.md
• https://huggingface.co/datasets/allenai/satlas-pretrain

The original dataset is not included in this model repository. The provided scripts use synthetic NPZ data to validate the full training, inference, evaluation, and visualization pipeline. This NPZ protocol is only a compact pipeline check and is not the original SatlasPretrain format.

Generate synthetic data:

python scripts/fake_data.py

The generated data/fake_satlaspretrain.npz contains four samples with images in [B, T, C, H, W] layout and labels for all seven task heads. Verify that conf/config.yaml points to this file before training.

Training

Single GPU or CPU:

python scripts/train.py

Multi-GPU:

torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py

Training outputs:

data/checkpoints/best.pth
data/checkpoints/last.pth

best.pth stores the checkpoint with the lowest joint multi-task validation loss, while last.pth stores the final training epoch. Checkpoints generated from synthetic data validate the pipeline only and are not pretrained weights reproducing the paper results.

Training Weights

This repository does not include synthetic, trained, or official pretrained weights. Running the training script generates best.pth and last.pth under data/checkpoints/. For real downstream use, train with appropriately converted SatlasPretrain data or use official AllenAI foundation model weights under their applicable license.

Inference

python scripts/inference.py

Prediction outputs:

result/output/segmentation.npy
result/output/regression.npy
result/output/point.npy
result/output/polygon.npy
result/output/polyline.npy
result/output/property.npy
result/output/classification.npy

These files store class probabilities, dense regression values, confidence maps, and category probabilities for the seven compact task heads.

Evaluation and Visualization

python scripts/result.py

Evaluation outputs:

result/output/satlaspretrain_metrics.json
result/output/qualitative_comparison.ppm

The metrics follow the task families reported in the paper: segmentation F1, regression MAE, point mAP, polygon mAP, polyline GEO, property accuracy, and classification accuracy. The qualitative comparison displays input imagery and paired ground-truth/prediction panels. Metrics and figures produced from synthetic data are pipeline checks, not paper results.

Official OneScience Resources

Platform	OneScience Main Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

• This repository is a reproduction of the original SatlasPretrain paper.
• The reproduction code is released under Apache License 2.0. The original SatlasPretrain data and official model weights remain subject to their respective source licenses.