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 or download them from GitHub:

pip install terratorch==1.2.1
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.

You can serve the model via vLLM with config.json.

Citation

Our technical report is released soon!