TerraTorch
blumenstiel's picture
Update ReadMe
abab3d3
metadata
license: apache-2.0
library_name: terratorch
datasets:
  - ibm-esa-geospatial/ImpactMesh-Flood
base_model:
  - ibm-esa-geospatial/TerraMind-1.0-base

arXiv Code IBMblog

TerraMind-base-Flood

TerraMind-base-Flood is based on TerraMind-1.0-base and was fine-tuned on ImpactMesh-Flood using TerraTorch. We use the Temporal Wrapper for a mid-fusion approach. The backbone processes the multimodal input while the decoder fuses the multi-temporal information. We refer to our technical report for details (coming soon!).

Usage

Quickstart with installing TerraTorch and the ImpactMesh DataModules:

pip install git+https://github.com/terrastackai/terratorch.git@multimodal
pip install impactmesh

Load the model via Lightning:

import torch
from terratorch.cli_tools import LightningInferenceModel

# Load TerraTorch task from 
task = LightningInferenceModel.from_config(
    "terramind_v1_base_impactmesh_flood.yaml",
    "TerraMind_v1_base_ImpactMesh_flood.pt",
)

model = task.model.model  # Get model from Lighting task
model.eval()

# Inputs with shape [B, C, T, H, W]
input = {
    "S2L2A": torch.randn([1, 12, 4, 256, 256]),
    "S1RTC": torch.randn([1, 2, 4, 256, 256]),
    "DEM": torch.randn([1, 1, 4, 256, 256]),  # Repeated per timestep
}

# Run inference
with torch.no_grad():
    pred = model(input).output

y_hat = pred.argmax(dim=1)

Run predictions via the TerraTorch CLI:

terratorch predict -c "terramind_v1_base_impactmesh_flood.yaml" --ckpt "TerraMind_v1_base_ImpactMesh_flood.pt" --predict_output_dir output/impactmesh_flood_predictions --predict_data_root "path/to/data/"

For prediction, the ImpactMesh data module expects a format similar to the training data with subfolders per modality and zarr.zip and tif files. Alternatively, you can adapt this inference code.

Citation

Our technical report is released soon!