This repository shows three examples of how Prithvi can be finetuned for downstream tasks. The examples include flood detection using Sentinel-2 data from the Sen1Floods11 dataset, burn scars detection using the NASA HLS fire scars dataset and multi-temporal crop classification using the NASA HLS multi-temporal crop classification dataset.
To finetune for these tasks in this repository, we make use of MMSegmentation, which provides an extensible framework for segmentation tasks.
MMSegmentation allows us to concatenate necks and heads appropriate for any segmentation downstream task to the encoder, and then perform the finetuning. This only requires setting up a config file detailing the desired model architecture, dataset setup and training strategy.
We build extensions on top of MMSegmentation to support our encoder and provide classes to read and augment remote sensing data (from .tiff files) using MMSegmentation data pipelines. These extensions can be found in the geospatial_fm directory, and they are installed as a package on the top of MMSegmentation for ease of use. If more advanced functionality is necessary, it should be added there.
The pretrained model we work with is a ViToperating as a Masked Autoencoder, trained on HLS data. The encoder from this model is made available as the backbone and the weights can be downloaded from Hugging Face here.
We use a simple architecture that adds a neck and segmentation head to the backbone. The neck concatenates and processes the transformer's token based embeddings into an embedding that can be fed into convolutional layers. The head processes this embedding into a segmentation mask. The code for the architecture can be found in this file.
Additionally, we provide extra components for data loading pipelines in geospatial_pipelines.py. These are documented in the file.
We observe the MMCV convention that all operations assume a channel-last format.
However, we also introduce some components with the prefix Torch, such as TorchNormalize. These components assume the torch convention of channel-first.
At some point during the pipeline, before feeding the data to the model, it is necessary to change to channel-first format.
We reccomend implementing the change after the ToTensor operation (which is also necessary at some point), using the TorchPermute operation.
Check out the exploration notebook for a more in depth example of the usage of the model.
conda create -n <environment-name> python==3.9conda activate <environment-name>pip install torch==1.11.0+cu115 torchvision==0.12.0+cu115 --extra-index-url https://download.pytorch.org/whl/cu115cd into the cloned repopip install -e .pip install -U openmimmim install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/{cuda_version}/{torch_version}/index.html. Note that pre-built wheels (fast installs without needing to build) only exist for some versions of torch and CUDA. Check compatibilities here: https://mmcv.readthedocs.io/en/v1.6.2/get_started/installation.html
mim install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/cu115/torch1.11.0/index.htmlconda create -n <environment-name> python=3.9conda activate <environment-name>pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113 git clone https://github.com/NASA-IMPACT/hls-foundation-os.git <your-local-path>\hls-foundation-osgit clone https://github.com/open-mmlab/mmsegmentation.git <your-local-path>\mmsegmentation cd <your-local-path>\mmsegmentation git checkout 186572a3ce64ac9b6b37e66d58c76515000c3280mmsegmentation @ git+https://github.com/open-mmlab/mmsegmentation.git@186572a3ce64ac9b6b37e66d58c76515000c3280 to mmsegmentation @ file:///<your-local-path>/mmsegmentationcd <your-local-path>\hls-foundation-ospip install -e .pip install -U openmimmim install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/{cuda_version}/{torch_version}/index.html. Note that pre-built wheels (fast installs without needing to build) only exist for some versions of torch and CUDA. Check compatibilities here: https://mmcv.readthedocs.io/en/v1.6.2/get_started/installation.html
mim install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/cu115/torch1.11.0/index.htmlconda install -c conda-forge opencvpip install datasets The flood detection dataset can be downloaded from Sen1Floods11. Splits in the mmsegmentation format are available in the data_splits folders.
The NASA HLS fire scars dataset can be downloaded from Hugging Face.
The NASA HLS multi-temporal crop classification dataset can be downloaded from Hugging Face.
Using git-lfs you can download the data as in the following example:
# from: https://huggingface.co/datasets/ibm-nasa-geospatial/multi-temporal-crop-classification
# Make sure you have git-lfs installed (https://git-lfs.com)
git lfs install
git clone https://huggingface.co/datasets/ibm-nasa-geospatial/multi-temporal-crop-classification
# extract files
cd multi-temporal-crop-classification
tar -xvf training_chips.tgz && tar -xvf validation_chips.tgzWithout git-lfs (Credit @robmarkcole):
mkdir data
cd data
mkdir multi-temporal-crop-classification
cd multi-temporal-crop-classification
# not this can take some time and appear to hang, be patient
wget https://huggingface.co/datasets/ibm-nasa-geospatial/multi-temporal-crop-classification/resolve/main/training_chips.tgz?download=true -O training_chips.tgz
tar -xvzf training_chips.tgz
wget https://huggingface.co/datasets/ibm-nasa-geospatial/multi-temporal-crop-classification/resolve/main/validation_chips.tgz?download=true -O validation_chips.tgz
tar -xvzf validation_chips.tgz
# delete some mac-os added files
find . -name '._*' -delete
# the following are NOT required (TBC)
https://huggingface.co/datasets/ibm-nasa-geospatial/multi-temporal-crop-classification/resolve/main/training_data.txt
https://huggingface.co/datasets/ibm-nasa-geospatial/multi-temporal-crop-classification/resolve/main/validation_data.txt
# instead copy over the files from the splits directory to the location of the images
cd ..
mkdir hls_burn_scars
cd hls_burn_scars
wget https://huggingface.co/datasets/ibm-nasa-geospatial/hls_burn_scars/resolve/main/hls_burn_scars.tar.gz?download=true -O hls_burn_scars.tar.gz
tar -xvf hls_burn_scars.tar.gzIn the configs folder there are three config examples for the three segmentation tasks. Complete the configs with your setup specifications. Parts that must be completed are marked with #TO BE DEFINED BY USER. They relate to the location where you downloaded the dataset, pretrained model weights, the test set (e.g. regular one or Bolivia out of bag data) and where you are going to save the experiment outputs.
a. With the conda env created above activated, run:
mim train mmsegmentation configs/sen1floods11_config.py or
mim train mmsegmentation configs/burn_scars.py or
mim train mmsegmentation configs/multi_temporal_crop_classification.py
b. Multi-gpu training can be run by adding --launcher pytorch --gpus <number of gpus>
c. To run testing:
mim test mmsegmentation configs/sen1floods11_config.py --checkpoint /path/to/best/checkpoint/model.pth --eval "mIoU" or
mim test mmsegmentation configs/burn_scars.py --checkpoint /path/to/best/checkpoint/model.pth --eval "mIoU" or
mim test mmsegmentation configs/multi_temporal_crop_classification.py --checkpoint /path/to/best/checkpoint/model.pth --eval "mIoU"
We also provide checkpoints on Hugging Face for the burn scars detection and the multi temporal crop classification tasks.
We provide a script to run inference on new data in GeoTIFF format. The data can be of any shape (e.g. height and width) as long as it follows the bands/channels of the original dataset. An example is shown below.
python model_inference.py -config /path/to/config/config.py -ckpt /path/to/checkpoint/checkpoint.pth -input /input/folder/ -output /output/folder/ -input_type tif -bands 0 1 2 3 4 5The bands parameter is useful in case the files used to run inference have the data in different orders/indexes than the original dataset.
This project builds on MMSegmentation and MMCV. For additional documentation, consult their docs (please note this is currently version 0.30.0 of MMSegmentation and version 1.5.0 of MMCV, not latest).
If this repository helped your research, please cite HLS foundation in your publications. Here is an example BibTeX entry:
@software{HLS_Foundation_2023,
author = {Jakubik, Johannes and Chu, Linsong and Fraccaro, Paolo and Bangalore, Ranjini and Lambhate, Devyani and Das, Kamal and Oliveira Borges, Dario and Kimura, Daiki and Simumba, Naomi and Szwarcman, Daniela and Muszynski, Michal and Weldemariam, Kommy and Zadrozny, Bianca and Ganti, Raghu and Costa, Carlos and Watson, Campbell and Mukkavilli, Karthik and Roy, Sujit and Phillips, Christopher and Ankur, Kumar and Ramasubramanian, Muthukumaran and Gurung, Iksha and Leong, Wei Ji and Avery, Ryan and Ramachandran, Rahul and Maskey, Manil and Olofossen, Pontus and Fancher, Elizabeth and Lee, Tsengdar and Murphy, Kevin and Duffy, Dan and Little, Mike and Alemohammad, Hamed and Cecil, Michael and Li, Steve and Khallaghi, Sam and Godwin, Denys and Ahmadi, Maryam and Kordi, Fatemeh and Saux, Bertrand and Pastick, Neal and Doucette, Peter and Fleckenstein, Rylie and Luanga, Dalton and Corvin, Alex and Granger, Erwan},
doi = {10.57967/hf/0952},
month = aug,
title = {{HLS Foundation}},
repository-code = {https://github.com/nasa-impact/hls-foundation-os},
year = {2023}
}