DOCSEARCHbyAlgolia
Submit Docs

Prithvi-EO-2.0 Foundation Model for Earth Observation

repository·main·Indexed 42 days ago

https://github.com/nasa-impact/prithvi-eo-2.0

A multi-temporal foundation model developed by IBM and NASA for Earth Observation (EO). Trained on 4.2M global time series samples of NASA Harmonized Landsat and Sentinel data at 30m resolution, it provides spatiotemporal embeddings for downstream tasks including flood detection, wildfire monitoring, crop segmentation, and CO2 flux regression. Available in 300M and 600M parameter versions, with optional temporal and location embeddings (-TL). Fine-tuning is supported via the TerraTorch library.

Tokens
4.3K
Snippets
12
Records
19
Agent score
56%

What's inside Prithvi-EO-2.0

  1. Overview of Prithvi-EO-2.0

    main

    Prithvi-EO-2.0 is a multi-temporal foundation model for Earth Observation (EO) applications. It is a more powerful successor to the original Prithvi model, trained on 4.2M global time series samples using NASA’s Harmonized Landsat and Sentinel data at 30m resolution. The model is designed to handle spatiotemporal characteristics and can incorporate geolocation and temporal metadata (date of acquisition) during pretraining and fine-tuning.

    Key features include:

    • Spatiotemporal Support: Uses 3D patch and positional embeddings to process sequences of T images of size (H, W).
    • Metadata Integration: Supports geolocation (latitude/longitude) and temporal (year/day-of-year) information via 2D sin/cos encoding.
    • Robustness: Includes a drop mechanism for metadata during pretraining to ensure the model performs even when geolocation or temporal data is missing.

  2. How the CO_2 Flux regression task works

    main

    The CO_2 Flux task is a regression model that processes HLS images and MERRA-2 data in parallel.

    • HLS Images: Processed through the prithvi_pretrained model.
    • MERRA-2 Data: Not passed through the Prithvi model directly.
    • Fusion: Both the Prithvi output (from HLS) and the MERRA-2 data are projected into the same embedding space to be combined for the final CO_2 flux prediction.

  5. Fine-tune Prithvi-EO-2.0 for downstream tasks

    main

    Fine-tuning is performed using the TerraTorch library. The repository provides sample configuration files (.yaml) and notebooks for several downstream Earth Observation tasks.

    Available Tasks and Configurations

    TaskDatasetTerraTorch Config/Code
    Flood DetectionSen1Floods11sen1floods11.yaml
    Wildfire Scar Detectionhls_burn_scarsfirescars.yaml
    Burn Scar Intensityburn_intensityburnintensity.yaml
    Landslide DetectionLandslide4senselandslide.yaml
    Multi-temporal Crop Segmentation (US)multi-temporal-crop-classificationmulticrop.yaml
    Multi-temporal Land Cover/Crop (Europe)sen4mapsen4map_land-cover.yaml, sen4map_crops.yaml
    Above Ground Biomass EstimationBioMasstersbiomassters.yaml

  6. Download the HLS and MERRA-2 dataset

    main

    The dataset for the CO_2 Flux regression task is available on Hugging Face (ibm-nasa-geospatial/hls_merra2_gppFlux). Use the provided download_dataset.py script to download it to a local directory. The script will automatically create train and test subdirectories within your specified save directory.

    python download_dataset.py --save_dir <directory to save the dataset>

  7. Configure dataset paths in fluxconfig_trainer.yaml

    main

    After downloading the data, you must update the fluxconfig_trainer.yaml configuration file to point to your local dataset paths. Specifically, update the data.chips and data.test_chips keys to match your directory structure.

    data:
  n_frame: 1
  chips: "/path/to/your/train/images/" # Update this path
  test_chips: "/path/to/your/test/images/" # Update this path
  input_size: [6,50, 50]
  means_for2018test: [0.07286696773903256, 0.10036772476940378, 0.11363777043869523, 0.2720510638470194, 0.2201167122609674, 0.1484162876040495]
  stds_for2018test: [0.13271414936598172, 0.13268933338964875, 0.1384673725283858, 0.12089142598551804, 0.10977084890500641, 0.0978705241034744]
...

  8. Execute Training and Inference with PyTorch Lightning

    main

    Training is managed by the PyTorch Lightning Trainer. You must define a PixelwiseRegressionTask (from torchgeo.trainers) which encapsulates the model, loss function, and optimizer.

    Training Steps:

    1. Define the task using PixelwiseRegressionTask(..., model=model_comb, loss="mse", optimizer="AdamW").
    2. Initialize the Trainer with desired callbacks (e.g., ModelCheckpoint, LearningRateMonitor).
    3. Call trainer.fit(model=task, datamodule=datamodule) to start training.
    4. Call trainer.predict(model=task, datamodule=datamodule) for inference.

    Inference/Prediction: trainer.predict returns a list of predictions. For the carbon flux task, you typically extract the first element of each result and concatenate them to evaluate against the ground truth.

    # Training
task = PixelwiseRegressionTask(None, None, model=model_comb, loss="mse", optimizer="AdamW")
trainer.fit(model=task, datamodule=datamodule)

# Inference
results = trainer.predict(model=task, datamodule=datamodule, return_predictions=True)

  9. Fine-tune Prithvi-EO-2.0 for Carbon Flux Regression

    main

    This notebook demonstrates how to fine-tune the Prithvi-EO-2.0 foundation model (specifically the 300M-TL version with temporal and location embeddings) to regress Gross Primary Production (GPP) flux using HLS (S30/L30) and MERRA-2 data.

    Workflow Overview:

    1. Setup: Download pre-trained weights from HuggingFace.
    2. Model Construction: Wrap the PrithviViT backbone in a custom prithvi_terratorch module and combine it with a RegressionModel_flux head.
    3. Data Preparation: Load HLS and MERRA-2 data using custom datamodules (flux_dataset and flux_dataloader).
    4. Training: Use PyTorch Lightning's Trainer with a PixelwiseRegressionTask to perform fine-tuning.
    5. Evaluation: Perform both 'zeroshot' (no training) and post-training evaluation, unnormalizing predictions to compute real-world metrics like $R^2$, MSE, and MAE.
    # To run the notebook
jupyter lab main_flux_finetune_baselines_trainer.ipynb

  10. Download Prithvi-EO-2.0 Checkpoints from HuggingFace

    main

    Before training, you must download the pre-trained weights. The example uses hf_hub_download to fetch the Prithvi_EO_V2_300M_TL.pt file from the ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL repository.

    from huggingface_hub import hf_hub_download
import os

if not os.path.isfile("Prithvi_EO_V2_300M_TL.pt"):
    hf_hub_download(
        repo_id="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL",
        filename="Prithvi_EO_V2_300M_TL.pt",
        local_dir=".",
    )

  12. Download the Landslide4Sense dataset

    main

    The Landslide4Sense dataset can be downloaded from Hugging Face using snapshot_download. Ensure you specify a DATASET_PATH for the local directory.

    from huggingface_hub import snapshot_download

DATASET_PATH = "./data"
repo_id = "ibm-nasa-geospatial/Landslide4sense"
_ = snapshot_download(repo_id=repo_id, repo_type="dataset", cache_dir="./cache", local_dir=DATASET_PATH)