README.md

---
license: apache-2.0
language:
- en
tags:
- OneScience
- Earth Science
- Remote Sensing Multi-Task Learning
- Multi-Temporal Remote Sensing
frameworks: PyTorch
datasets:
- SatlasPretrain
- NAIP
- Sentinel-2
---

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

# Model Introduction

SatlasPretrain is a multi-task pretraining model for large-scale remote sensing understanding. It uses a dual-stream hierarchical backbone to process multi-temporal high-resolution RGB and Sentinel-2 multispectral imagery and jointly learns dense prediction and global classification tasks.

Paper: SatlasPretrain: A Large-Scale Dataset for Remote Sensing Image Understanding  
https://arxiv.org/abs/2211.15660

# Model Description

SatlasPretrain was proposed by a research team from the Allen Institute for AI. The model is trained with NAIP-style high-resolution imagery, Sentinel-2 multispectral imagery, and multiple types of remote sensing labels from SatlasPretrain. It is suitable for remote sensing multi-task learning, including semantic segmentation, regression, point, polyline, and polygon object prediction, property recognition, and scene classification.

# Use Cases

| Scenario | Description |
| :---: | :--- |
| Multi-temporal remote sensing fusion | Jointly process 4 time steps of high-resolution RGB and 8 time steps of Sentinel-2 data. |
| Remote sensing multi-task learning | Simultaneously predict segmentation, regression, point, polygon, polyline, property, and classification tasks. |
| Land-cover and crop mapping | Use multi-temporal imagery for semantic segmentation tasks such as land cover and crop type mapping. |
| Infrastructure feature extraction | Identify point, polyline, and polygon geographic objects such as buildings, roads, railways, airports, and energy facilities. |
| Local engineering validation | Use a small amount of synthetic data to check the training, inference, and evaluation workflows. |
| Multi-GPU training | Launch distributed training with `torchrun`. |

# Usage Guide

## 1. OneCode

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

[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Download and Installation

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

### Environment Dependencies

**Hardware Requirements**

- A GPU or DCU is recommended.
- CPU can be used for small-configuration connectivity validation; 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 current cluster, is recommended.

**DCU Environment**

```bash
# 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**

```bash
# 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

By default, 1 training and 1 test synthetic sample are used to validate the engineering workflow. They are saved as `data/train.npz` and `data/test.npz`, respectively.

The synthetic data preserves the official multi-temporal configuration of 4 NAIP time steps, 8 Sentinel-2 time steps, a `512x512` spatial size, and seven task-label specifications.

Real data must be preprocessed and converted to the following NPZ training protocol. This protocol is consistent with the model input specification but is not the download format of the original datasets.

```text
highres_images: float32 [N,4,3,512,512]
lowres_images: float32 [N,8,9,512,512]
valid_highres_times: bool [N,4]
valid_lowres_times: bool [N,8]
sample_ids: string [N]
segmentation: int64 [N,512,512]
regression: float32 [N,1,512,512]
point: float32 [N,1,512,512]
polygon: float32 [N,1,512,512]
polyline: float32 [N,1,512,512]
property: int64 [N]
classification: int64 [N]
```

`fake_data.py` automatically writes the `protocol` and `source` protocol metadata. These fields must be retained when using real data.

```bash
python scripts/fake_data.py
```

### Training

```bash
python scripts/train.py
```

For multi-GPU training, use:

```bash
torchrun --nproc_per_node=8 scripts/train.py
```

Training jointly optimizes seven remote sensing task types and saves a checkpoint and aggregate training metrics. The default configuration is intended for quick workflow validation. Formal experiments should use the multi-temporal data, complete task labels, model configuration, and training duration corresponding to the paper.

```text
result/checkpoints/satlaspretrain.pt
result/training/metrics.json
```

### Training Weights

This repository will provide SatlasPretrain training weights in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.

### Inference

```bash
python scripts/inference.py
```

Inference loads the training checkpoint, generates predictions for seven task types, retains sample identity and protocol metadata, and saves the results to:

```text
result/output/predictions.npz
```

### Evaluation and Visualization

```bash
python scripts/result.py
```

Evaluation covers segmentation, regression, point, polyline, and polygon objects, property, and classification tasks, and generates a multi-task prediction figure. Results on synthetic data are only for engineering workflow validation and do not represent full-paper performance.

```text
result/evaluation/metrics.json
result/evaluation/multitask_predictions.png
```

# 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.

Use of the code and data in this repository remains subject to the licenses and terms of use of their respective projects.