Image segmentation by foundation model finetuning

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.

The approach

Background

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 backbone

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.

The architectures

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.

The pipeline

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.

Tutorial

Check out the exploration notebook for a more in depth example of the usage of the model.

Setup

Dependencies

1. Clone this repository
2. conda create -n <environment-name> python==3.9
3. conda activate <environment-name>
4. Install torch (tested for >=1.7.1 and <=1.11.0) and torchvision (tested for >=0.8.2 and <=0.12). May vary with your system. Please check at: https://pytorch.org/get-started/previous-versions/.
   1. e.g.: pip install torch==1.11.0+cu115 torchvision==0.12.0+cu115 --extra-index-url https://download.pytorch.org/whl/cu115
5. cd into the cloned repo
6. pip install -e .
7. pip install -U openmim
8. mim 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
   1. e.g.: mim install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/cu115/torch1.11.0/index.html

Alternate Setup (Windows Users - Tested for Windows 10)

1. conda create -n <environment-name> python=3.9
2. conda activate <environment-name>
3. Install torch (tested for >=1.7.1 and <=1.11.0) and torchvision (tested for >=0.8.2 and <=0.12). May vary with your system. Please check at: https://pytorch.org/get-started/previous-versions/.
   1. e.g.: 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
4. git clone https://github.com/NASA-IMPACT/hls-foundation-os.git <your-local-path>\hls-foundation-os
5. git clone https://github.com/open-mmlab/mmsegmentation.git <your-local-path>\mmsegmentation
6. cd <your-local-path>\mmsegmentation
7. Checkout mmsegmentation version compatible with hls-foundation: git checkout 186572a3ce64ac9b6b37e66d58c76515000c3280
8. modify setup.py so it installs from the cloned mmsegmentation. Change line mmsegmentation @ git+https://github.com/open-mmlab/mmsegmentation.git@186572a3ce64ac9b6b37e66d58c76515000c3280 to mmsegmentation @ file:///<your-local-path>/mmsegmentation
9. cd <your-local-path>\hls-foundation-os
10. pip install -e .
11. pip install -U openmim
12. mim 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
    1. e.g.: mim install mmcv-full==1.6.2 -f https://download.openmmlab.com/mmcv/dist/cu115/torch1.11.0/index.html
13. conda install -c conda-forge opencv
14. pip install datasets

Data

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

Without 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.gz

Running the finetuning

1. In 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.

2. 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"

Checkpoints on Hugging Face

We also provide checkpoints on Hugging Face for the burn scars detection and the multi temporal crop classification tasks.

Running the inference

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 5

The bands parameter is useful in case the files used to run inference have the data in different orders/indexes than the original dataset.

Additional documentation

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

Citation

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}
}