Home

Awesome

Multi-modal learning for geospatial vegetation forecasting

Code for Benson et. al., CVPR (2024)

arXiv DOI

:zap: Our paper on Multi-modal learning for geospatial vegetation forecasting has been accepted to CVPR! We benchmark a wide range of EarthNet models on the new GreenEarthNet dataset, plus introducing a new transformer-based SOTA: Contextformer.

Installation

The easiest is to just create a conda environment as follows:

conda create -n greenearthnet python=3.10
conda activate greenearthnet
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia
conda install -c conda-forge u8darts-notorch
conda install -c conda-forge numpy matplotlib pillow xarray zarr netcdf4 scipy imageio networkx scikit-image s3fs pyproj pyarrow seaborn cartopy tensorboard dask pytorch-lightning=1.7.7 torchmetrics=0.7.3 statsmodels
pip install earthnet earthnet-minicuber segmentation-models-pytorch albumentations
pip install git+https://github.com/earthnet2021/earthnet-models-pytorch.git@v0.1.0

GreenEarthNet Dataset Download

We use the EarthNet Toolkit to download the GreenEarthNet dataset. First make sure you have it installed (pip install earthnet) and also make sure you have enough free disk space! We recommend 1TB. Then, downloading the dataset or a part thereof is as simple as:

import earthnet as entk
entk.download(dataset = "greenearthnet", split = "train", save_directory = "data_dir")

Where data_dir is the directory where EarthNet2021 shall be saved and split is "all"or a subset of ["train","iid","ood","extreme","seasonal"].

NOTE: In the paper the dataset is called GreenEarthNet, but in the codebase you will also find the following acronyms that were used during development: earthnet2021x, en21x.

Model Inference

For training and inference, we are using earthnet-models-pytorch, which means this codebase only contains a small python script and then a lot of config files, alongside with the pre-trained weights stored on Zenodo (see below).

python test.py model_configs/path/to/config.yaml weights/path/to/weights.ckpt --track ood-t_chopped --pred_dir preds/path/to/save/predictions --data_dir data/path/to/dataset

You may also evaluate the model predictions and compute the metrics reported in the paper:

python eval.py path/to/test/dataset path/to/model/predictions path/for/saving/scores

You can optionally include a comparison directory, to also compute the outperformance score, e.g. relative to the Climatology baseline as done in the paper: --compare_dir path/to/climatology/predictions

Model Training

python train.py model_configs/path/to/config.yaml --data_dir data/path/to/dataset

Downloading Pre-Trained Model Weights

The pre-trained model weights are stored on Zenodo. You may download them (2.3GB) using this code:

wget https://zenodo.org/records/10793870/files/model_weights.zip
unzip model_weights.zip

DOI

Dataset Generation

The GreenEarthNet dataset has been generated with the EarthNet minicuber (pip install earthnet-minicuber). You can create more minicubes in the same style using the following code snippet:

import earthnet_minicuber as emc

specs = {
    "lon_lat": (43.598946, 3.087414),  # center pixel
    "xy_shape": (128, 128),  # width, height of cutout around center pixel
    "resolution": 20,  # in meters.. will use this together with grid of primary provider..
    "time_interval": "2018-01-01/2021-12-31",
    "providers": [
        {
            "name": "s2",
            "kwargs": {
                "bands": ["B02", "B03", "B04", "B8A", "SCL"],  # , "B09", "B11", "B12"],
                "best_orbit_filter": False,
                "five_daily_filter": True,
                "brdf_correction": False,
                "cloud_mask": True,
                "cloud_mask_rescale_factor": 2,
                "aws_bucket": "planetary_computer",
            },
        },
        {"name": "srtm", "kwargs": {"bands": ["dem"]}},
        {"name": "alos", "kwargs": {}},
        {"name": "cop", "kwargs": {}},
        {
            "name": "esawc",
            "kwargs": {"bands": ["lc"], "aws_bucket": "planetary_computer"},
        },
        # {
        #     "name": "geom",
        #     "kwargs": {"filepath": "downloads/Geomorphons/geom/geom_90M_africa_europe.tif"}
        # }
        # Also Missing here: EOBS v26 https://surfobs.climate.copernicus.eu/dataaccess/access_eobs.php#datafiles
    ],
}

mc = emc.load_minicube(specs, compute = True)

To fully obtain minicubes aligned with the GreenEarthNet dataset that are useful for prediction, you still need to download and merge E-OBS v26.

Cloud Mask Training

The Cloud Mask Algorithm is trained on the CloudSEN12 dataset, if you want to repeat the training, follow the instructions on their website to download the dataset and then run the python script:

python cloudmask/train_cloudmask_l2argbnir.py

NOTE: This python script is based upon the original CloudSEN12 MobileNetv2 implementation!

If you wish to just use the trained Cloud Mask Algorithm, you may do so, it is included in the EarthNet Minicuber package (pip install earthnet-minicuber).

Reproducing Paper Tables

Table 2

ModelConfig/Weights
Persistencepython model_pixelwise/persistence.py --data_dir data/path/to/dataset --out_dir experiments/path/to/save/predictions
Previous Yearpython model_pixelwise/previousyear.py --data_dir data/path/to/dataset --out_dir experiments/path/to/save/predictions 
Climatologypython model_pixelwise/climatology.py --data_dir data/path/to/dataset --out_dir experiments/path/to/save/predictions 
Kalman Filter python model_pixelwise/predict_kalman_xgb_prophet.py minicube_idx --data_dir data/path/to/dataset --out_dir experiments/path/to/save/predictions 
 LightGBMpython model_pixelwise/predict_kalman_xgb_prophet.py minicube_idx --data_dir data/path/to/dataset --out_dir experiments/path/to/save/predictions 
 Prophet python model_pixelwise/predict_kalman_xgb_prophet.py minicube_idx --data_dir data/path/to/dataset --out_dir experiments/path/to/save/predictions 
Diaconu ConvLSTMSee Weather2Land Github
Kladny SG-ConvLSTMSee SatelliteImageForecasting Github
Earthformer EarthNet2021 weightsSee Earthformer Github 
ConvLSTM Seed <42, 97, 27>model_configs/convlstm/convlstm1M/seed=<42,97,27>.<yaml/ckpt>
Earthformer retrainedSee Earthformer Github
PredRNN Seed <42, 97, 27>model_configs/predrnn/predrnn1M/seed=<42,97,27>.<yaml/ckpt>
SimVP Seed <42, 97, 27>model_configs/simvp/simvp6M/seed=<42,97,27>.<yaml/ckpt>
Contextformer Seed <42, 97, 27>model_configs/contextformer/contextformer6M/seed=<42,97,27>.<yaml/ckpt>

Table 3

ModelConfig/Weights
Climatology
1x1 LSTMmodel_configs/lstm1x1/<lstm1x1,spatialshuffle>/seed=42.<yaml/ckpt>
Next-frame UNetmodel_configs/nfunet/<nfunet,spatialshuffle>/seed=42.<yaml/ckpt>
Next-cuboid UNet model_configs/nfunet/<nfunet,spatialshuffle>/seed=42.<yaml/ckpt>
ConvLSTMmodel_configs/convlstm/<convlstm1M,spatialshuffle>/seed=42.<yaml/ckpt>
PredRNNmodel_configs/predrnn/<predrnn1M,spatialshuffle>/seed=42.<yaml/ckpt>
SimVPmodel_configs/simvp/<simvp6M,spatialshuffle>/seed=42.<yaml/ckpt>
Contextformermodel_configs/contextformer/<contextformer6M,spatialshuffle>/seed=42.<yaml/ckpt>

Table 4

ModelConfig/Weights
MLP vision encodermodel_configs/contextformer/mlp_vision_enc/seed=42.<yaml/ckpt>
PVT encoder (frozen)model_configs/contextformer/pvt_vision_frozen/seed=42.<yaml/ckpt>
w/ cloud mask tokenmodel_configs/contextformer/mask_clouds/seed=42.<yaml/ckpt>
w/ learned Vhat0model_configs/contextformer/learned_vhat0/seed=42.<yaml/ckpt>
w/ last pixel Vhat0model_configs/contextformer/last_vhat0/seed=42.<yaml/ckpt>
Contextformer 6Mmodel_configs/contextformer/contextformer6M/seed=<42,97,27>.<yaml/ckpt>
Contextformer 16Mmodel_configs/contextformer/contextformer16M/seed=<42,97,27>.<yaml/ckpt>

Table 5

ModelConfig/Weights
Climatology
Contextformer 6Mmodel_configs/contextformer/contextformer6M/seed=<42,97,27>.<yaml/ckpt>

Reproducing Paper Figures

Figure 3 and 5

ModelConfig/Weights
Contextformer 6Mmodel_configs/contextformer/contextformer6M/seed=42.<yaml/ckpt>

Figure 4

ModelConfig/Weights
ConvLSTMmodel_configs/convlstm/<convlstm1M,noweather>/seed=<42,97,27>.<yaml/ckpt>
PredRNNmodel_configs/predrnn/<predrnn1M,noweather>/seed=<42,97,27>.<yaml/ckpt>
SimVPmodel_configs/simvp/<simvp6M,noweather>/seed=<42,97,27>.<yaml/ckpt>
Contextformermodel_configs/contextformer/<contextformer6M,noweather>/seed=<42,97,27>.<yaml/ckpt>

Citation

@InProceedings{Benson_2024_CVPR,
    author    = {Benson, Vitus and Robin, Claire and Requena-Mesa, Christian and Alonso, Lazaro and Carvalhais, Nuno and Cort\'es, Jos\'e and Gao, Zhihan and Linscheid, Nora and Weynants, M\'elanie and Reichstein, Markus},
    title     = {Multi-modal Learning for Geospatial Vegetation Forecasting},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2024},
    pages     = {27788-27799}
}