Home

Awesome

Satlas Super Resolution

Satlas aims to provide open AI-generated geospatial data that is highly accurate, available globally, and updated on a frequent (monthly) basis. One of the data applications in Satlas is globally generated Super-Resolution imagery for 2023.

<p align="center"> <img src="figures/kenya_sentinel2.gif" alt="animated" width=300 height=300 /> <img src="figures/kenya_superres.gif" alt="animated" width=300 height=300 /> </p>

We describe the many findings that led to the global super-resolution outputs in the paper, Zooming Out on Zooming In: Advancing Super-Resolution for Remote Sensing. Supplementary material is available here.

<p align="center"> <img src="figures/teaser.svg" /> </p>

This repository contains the training and inference code for the AI-generated Super-Resolution data found at https://satlas.allen.ai/, as well as code, data, and model weights corresponding to the paper.

The experiments branch contains config files for experiments from the paper, while the main branch is limited to showcasing the main features.

Installation

Initialize conda:

conda create --name ssr python==3.9
conda activate ssr
pip install -r requirements.txt
conda install gdal

Download

Data

There are two training sets:

The urban set (termed S2-NAIP) was used for all experiments in the paper, because we found the full set to be overwhelmed with monotonous landscapes.

There are three val/test sets:

Additional data includes:

All of the above data (except for the full training set due to size) can be downloaded at this link, or individual links are provided above for ease of downloading.

HuggingFace

The train_urban_set, split into many partitions, val_set, and test_set are available for download on HuggingFace as well.

Model Weights

Weights from models trained on the S2-NAIP dataset are listed below.

ESRGAN

Varying number of input Sentinel-2 images (just RGB bands):

Number Input ImagesWeightsConfig
11-S2-imagesesrgan_baseline_1S2.yml
22-S2-imagesesrgan_baseline_2S2.yml
44-S2-imagesesrgan_baseline_4S2.yml
88-S2-imagesesrgan_baseline.yml
1616-S2-imagesesrgan_baseline_16S2.yml

Different Sentinel-2 bands used as input (8 input images):

BandsWeightsConfig
10m10m-S2-bandsesrgan_baseline_10m.yml
20m20m-S2-bandsesrgan_baseline_20m.yml
60m60m-S2-bandsesrgan_baseline_60m.yml

SRCNN & HighResNet

S2-NAIP Dataset Structure

The dataset consists of image pairs from Sentinel-2 and NAIP satellites, where a pair is a time series of Sentinel-2 images that overlap spatially and temporally [within 3 months] with a NAIP image. The imagery is from 2019-2020 and is limited to the USA.

<p align="center"> <img src="figures/image_pair.svg" /> </p>

The images adhere to the same Web-Mercator tile system as in SatlasPretrain.

NAIP

The NAIP images included in this dataset are 25% of the original NAIP resolution. Each image is 128x128px with RGB channels.

In each set, there is a naip folder containing images in this format: naip/{image_uuid}/{tile}/rgb.png, where image_uuid is the image's unique identifier with the capture timestamp, and tile refers to its location in a 2^17 x 2^17 Web-Mercator grid (ex. 12345_67890).

Sentinel-2

We use the Sentinel-2 L1C imagery. Models that input 3 bands use the TCI file provided by ESA. This contains an 8-bit image that has been normalized by ESA to the 0-255 range. The image is normalized for input to the model by dividing the 0-255 RGB values by 255, and retaining the RGB order. Most experiments utilize just TCI, but for non-TCI bands, the 16-bit source data is divided by 8160 and clipped to 0-1.

For each NAIP image, there is a time series of corresponding 32x32px Sentinel-2 images. These time series are saved as pngs in the shape, [number_sentinel2_images * 32, 32, 3]. Note that the input images do not need to be in chronological order.

In each set, there is a sentinel2 folder containing these time series in the format: sentinel2/{tile}/{band}.png, where tile refers to its location in a 2^17 x 2^17 Web-Mercator grid (ex. 12345_67890) and band refers to the Sentinel-2 bands (tci, b01, b05, b06, b07, b08, b09, b10, b11, b12).

How to Process Raw Sentinel-2 Data

For each image, use the TCI.jp2 file and use the following code to process it in the same way as the S2-NAIP dataset:

import rasterio

tci_jp2_path = 'path/to/TCI.jp2'
with rasterio.open(tci_jp2_path) as src:
    img_rep, meta_rep = reproject_image(
        img, meta, rasterio.crs.CRS.from_epsg(3857), resolution=(9.555, 9.555), resampling=rasterio.warp.Resampling.bilinear)

with rasterio.open(tci_jp2_path.replace('.jp2', '_rep.jp2', 'w', **meta_rep) as dst:
    dst.write(img_rep)

Model

In the paper, we experiment with SRCNN, HighResNet, SR3, and ESRGAN. For a good balance of output quality and inference speed, we use the ESRGAN model for generating global super-resolution outputs.

Our ESRGAN model is an adaptation of the original ESRGAN, with changes that allow the input to be a time series of Sentinel-2 images. All models are trained to upsample by a factor of 4.

<p align="center"> <img src="figures/esrgan_generator.svg" /> </p>

The SR3 diffusion model code has lived in a separate repository. We are working to release that as well.

Training

To train a model on this dataset, run the following command, with the desired configuration file:

python -m ssr.train -opt ssr/options/esrgan_s2naip_urban.yml

There are several sample configuration files in ssr/options/. Make sure the configuration file specifies correct paths to your downloaded data, the desired number of low-resolution input images, model parameters, and pretrained weights (if applicable).

Add the --debug flag to the above command if wandb logging, model saving, and visualization creation is not wanted.


To train with multiple GPUs, use the following command:

PYTHONPATH=. python -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 ssr/train.py -opt ssr/options/esrgan_s2naip_urban.yml --launcher pytorch

Testing

To evaluate the model on a validation or test set, when ground truth high-res images are available, run the following command, with the desired configuration file:

python -m ssr.test -opt ssr/options/esrgan_s2naip_urban.yml

This will test the model using data and parameters specified in ['datasets']['test'], and will save the model outputs as pngs in the results/ directory. Specified metrics will be displayed to the screen at the end.

Inference

To run inference on data, when ground truth high-res images are not available, run the following command:

python -m ssr.infer -opt ssr/options/infer_example.yml

Inference settings are specified in the configuration file. The data_dir can be of any directory structure, but must contain pngs. Both the original low-res images and the super-res images will be saved to the save_path.


When running inference on an entire Sentinel-2 tile (consisting of a 16x16 grid of chunks), there is the infer_grid.py script that will stitch the individual chunks together into one large image.

Try this out on the S2NAIP test set with this command:

python -m ssr.infer_grid -opt ssr/options/infer_grid_example.yml

<p align="center"> <img src="figures/stitch_example.svg" height=300 /> </p>

Accuracy

There are instances where the generated super resolution outputs are incorrect.

Specifically:

  1. Sometimes the model generates vessels in the water or cars on a highway, but because the input is a time series of Sentinel-2 imagery (which can span a few months), it is unlikely that those things persist in one location.
<p align="center"> <img src="figures/vessel_hallucination.jpg" width=300 height=300 /> <img src="figures/car_hallucination.jpg" width=300 height=300 /> </p>
  1. Sometimes the model generates natural objects like trees or bushes where there should be a building, or vice versa. This is more common in places that look vastly different from the USA, such as the example below in Kota, India.
<p align="center"> <img src="figures/kota_india.svg" height=300 /> </p>

Acknowledgements

Thanks to these codebases for foundational Super-Resolution code and inspiration:

BasicSR

Real-ESRGAN

Image Super-Resolution via Iterative Refinement (SR3)

WorldStrat

Citation

If you find this work useful, please consider citing our work:

@misc{satlassuperres,
      title={Zooming Out on Zooming In: Advancing Super-Resolution for Remote Sensing}, 
      author={Piper Wolters and Favyen Bastani and Aniruddha Kembhavi},
      year={2023},
      eprint={2311.18082},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2311.18082}, 
}

Contact

If you have any questions, please email piperw@allenai.org or open an issue.