OlmoEarth Projects

This repository contains configuration files, model checkpoint references, and
documentation for several remote sensing models built on top of OlmoEarth at Ai2. It
also includes tooling and tutorials for building new models using various components of
OlmoEarth.

The models available here are:

• Live Fuel Moisture Content Mapping
• Forest Loss Driver Classification
• Mangrove Mapping
• Ecosystem Type Mapping
• Land Use / Land Cover Mapping in Southern Kenya

The links above provide more details about the training data and intended use case for
each model.

Here are tutorials for applying OlmoEarth for new tasks:

• Fine-tuning OlmoEarth for Segmentation
• Computing Embeddings using OlmoEarth
• Fine-tuning OlmoEarth in rslearn

These tutorials use all or a subset of the components of OlmoEarth:

• olmoearth_pretrain, the OlmoEarth
  pre-trained model.
• rslearn, our tool for obtaining satellite
  images and other geospatial data from online data sources, and for fine-tuning
  remote sensing foundation models.
• olmoearth_run, our higher-level
  infrastructure that automates various steps on top of rslearn such as window creation
  and inference post-processing.

Installation

We recommend installing using uv. See
Installing uv for
instructions to install uv. Once uv is installed:

git clone https://github.com/allenai/olmoearth_projects.git
cd olmoearth_projects
uv sync
source .venv/bin/activate

Applying Existing Models

There are three steps to applying the models in this repository:

1. Customize the prediction request geometry, which specifies the spatial and temporal
   extent to run the model on.
2. Execute the olmoearth_run steps to build an rslearn dataset for inference, and to
   apply the model on the dataset.
3. Collect and visualize the outputs.

Customizing the Prediction Request Geometry

The configuration files for each project are stored under
olmoearth_run_data/PROJECT_NAME/. There are three configuration files:

• dataset.json: this is an rslearn dataset configuration file that specifies the
  types of satellite images that need to be downloaded to run the model, and how to
  obtain them. Most models rely on some combination of Sentinel-1 and Sentinel-2
  satellite images, and are configured to download those images from Microsoft
  Planetary Computer.
• model.yaml: this is an rslearn model configuration file that specifies the model
  architecture, fine-tuning hyperparameters, data loading steps, etc.
• olmoearth_run.yaml: this is an olmoearth_run configuration file that specifies how
  the prediction request geometry should be translated into rslearn windows, and how
  the inference outputs should be combined together.

Some projects also include an example prediction_request_geometry.geojson, but this
will need to be modified to specify your target region. The spatial extent is specified
with standard GeoJSON features; you can use geojson.io to draw
polygons on a map and get the corresponding GeoJSON. The temporal extent is specified
using properties on each feature:

{
  "type": "FeatureCollection",
  "properties": {},
  "features": [
    {
      "type": "Feature",
      "geometry": {
        // ...
      },
      "properties": {
        "oe_start_time": "2024-01-01T00:00:00+00:00",
        "oe_end_time": "2024-02-01T00:00:00+00:00"
      },
    }
  ]
}

Here, the oe_start_time and oe_end_time indicate that the prediction for the
location of this feature should be based on satellite images around January 2024. The
per-model documentation details how these timestamps should be chosen. Some models like
forest loss driver classification provide project-specific tooling for generating the
prediction request geometry.

Executing olmoearth_run

Consult the per-model documentation to download the associated fine-tuned model
checkpoint. For example:

mkdir ./checkpoints
wget https://huggingface.co/allenai/OlmoEarth-v1-FT-LFMC-Base/resolve/main/model.ckpt -O checkpoints/lfmc.ckpt

Set needed environment variables:

export NUM_WORKERS=32
export WANDB_PROJECT=lfmc
export WANDB_NAME=lfmc_inference_run
export WANDB_ENTITY=YOUR_WANDB_ENTITY

Then, execute olmoearth_run:

mkdir ./project_data
python -m olmoearth_projects.main olmoearth_run olmoearth_run --config_path $PWD/olmoearth_run_data/lfmc/ --checkpoint_path $PWD/checkpoints/lfmc.ckpt --scratch_path project_data/lfmc/

Visualizing Outputs

The results directory (project_data/lfmc/results/results_raster/ in the example)
should be populated with one or more GeoTIFFs. You can visualize this in GIS software
like qgis:

qgis project_data/lfmc/results/results_raster/*.tif

Reproducing Fine-tuning for Existing Models

We have released model checkpoints for each of the fine-tuned models in this
repository, but you can reproduce the model by fine-tuning the pre-trained OlmoEarth
checkpoint on each task training dataset.

First, consult the per-model documentation above for the URL of the rslearn dataset tar
file, and download and extract it. For example, for the LFMC model:

wget https://huggingface.co/datasets/allenai/olmoearth_projects_lfmc/blob/main/dataset.tar
tar xvf dataset.tar

Set environment variables expected by the fine-tuning procedure (uses W&B)

export DATASET_PATH=/path/to/extracted/data/
export NUM_WORKERS=32
export TRAINER_DATA_PATH=./trainer_data
export PREDICTION_OUTPUT_LAYER=output
export WANDB_PROJECT=olmoearth_projects
export WANDB_NAME=my_training_run
export WANDB_ENTITY=...

Then run fine-tuning using the model configuration file in the olmoearth_run_data,
e.g.:

rslearn model fit --config olmoearth_run_data/lfmc/model.yaml

Losses and metrics should then be logged to your W&B. The checkpoint would be saved in
the TRAINER_DATA_PATH (e.g. ./trainer_data); two checkpoints should be saved, the
latest checkpoint (last.ckpt) and the best checkpoint (epoch=....ckpt). You can use
the best checkpoint for the Applying Existing Models section in lieu of the checkpoint
that we proivde.

If training fails halfway, you can resume it from last.ckpt:

rslearn model fit --config olmoearth_run_data/lfmc/model.yaml --ckpt_path $TRAINER_DATA_PATH/last.ckpt

License

This code is licensed under the OlmoEarth Artifact License.