Update ReadMe

README.md

@@ -1,3 +1,68 @@
 ---
 license: apache-2.0
+library_name: terratorch
+datasets:
+  - ibm-esa-geospatial/ImpactMesh-Fire
+base_model:
+  - ibm-esa-geospatial/TerraMind-1.0-base
 ---
+
+[![arXiv](https://img.shields.io/badge/arXiv-comming_soon-b31b1b?logo=arxiv)](https://arxiv.org/abs/todo)
+[![Code](https://img.shields.io/badge/GitHub-ImpactMesh-EE4B2B?logo=github)](https://github.com/IBM/ImpactMesh)
+[![IBMblog](https://img.shields.io/badge/Blog-IBM-0F62FE)](https://research.ibm.com/blog/todo)
+
+# TerraMind-base-Fire
+
+TerraMind-base-Fire is based on [TerraMind-1.0-base](https://huggingface.co/ibm-esa-geospatial/TerraMind-1.0-base) and was fine-tuned on [ImpactMesh-Fire](https://huggingface.co/datasets/ibm-esa-geospatial/ImpactMesh-FLood) using [TerraTorch](https://terrastackai.github.io/terratorch/stable/).
+We use the [Temporal Wrapper](https://terrastackai.github.io/terratorch/stable/guide/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](https://github.com/IBM/ImpactMesh):
+```shell
+pip install git+https://github.com/terrastackai/terratorch.git@multimodal
+pip install impactmesh
+```
+
+Load the model via Lightning:
+
+```python
+import torch
+from terratorch.cli_tools import LightningInferenceModel
+
+# Load TerraTorch task from 
+task = LightningInferenceModel.from_config(
+    "terramind_v1_base_impactmesh_fire.yaml",
+    "TerraMind_v1_base_ImpactMesh_fire.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:
+
+```shell
+terratorch predict -c "terramind_v1_base_impactmesh_fire.yaml" --ckpt "TerraMind_v1_base_ImpactMesh_fire.pt" --predict_output_dir output/impactmesh_fire_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](https://github.com/IBM/ImpactMesh/blob/main/impactmesh/run_inference.py).
+
+## Citation
+
+Our technical report is released soon!

terramind_v1_base_impactmesh_fire.yaml

@@ -0,0 +1,123 @@
+# lightning.pytorch==2.1.1
+seed_everything: 42
+trainer:
+  accelerator: auto
+  strategy: auto
+  devices: auto
+  num_nodes: 1
+  precision: 16-mixed
+#  logger: true # Set to logger: true for Tensorboard
+  logger:
+    class_path: lightning.pytorch.loggers.WandbLogger
+    init_args:
+      project: ImpactMesh-Fire
+      name: TerraMind-Base_lr1e-4
+      config:
+        modalities: S1-S2-DEM
+        dataset: IM-Fire
+
+  callbacks:
+    - class_path: RichProgressBar
+    - class_path: LearningRateMonitor
+      init_args:
+        logging_interval: epoch
+    - class_path: EarlyStopping
+      init_args:
+        monitor: val/loss
+        patience: 10
+    - class_path: ModelCheckpoint
+      init_args:
+        monitor: val/loss
+        mode: min
+        save_weights_only: true
+        dirpath: output/terramind_base_impactmesh_fire/1e-4
+        filename: best_val_loss
+  max_epochs: 50
+  log_every_n_steps: 5
+  default_root_dir: output/terramind_base_impactmesh_fire/
+data:
+  class_path: impactmesh.impactmesh_datamodule.ImpactMeshDataModule
+  init_args:
+    batch_size: 16
+    num_workers: 8
+    data_root: data/ImpactMesh-Fire/data
+    train_split: data/ImpactMesh-Fire/split/impactmesh_wildfire_train.txt
+    val_split: data/ImpactMesh-Fire/split/impactmesh_wildfire_val.txt
+    test_split: data/ImpactMesh-Fire/split/impactmesh_wildfire_test.txt
+    timesteps: [0, 1, 2, 3]
+    modalities:
+        - S2L2A
+        - S1RTC
+        - DEM
+    no_data_replace: 0
+    train_transform:
+      - class_path: terratorch.datasets.transforms.FlattenTemporalIntoChannels
+      - class_path: albumentations.D4
+      - class_path: albumentations.pytorch.ToTensorV2
+      - class_path: terratorch.datasets.transforms.UnflattenTemporalFromChannels
+        init_args:
+          n_timesteps: 4
+# Use pretraining stats with frozen encoder
+#    means:
+#      S2L2A: [ 1390.458, 1503.317, 1718.197, 1853.91, 2199.1, 2779.975, 2987.011, 3083.234, 3132.22, 3162.988, 2424.884, 1857.648 ]
+#      S1RTC: [ -10.93, -17.329 ]
+#      DEM: [ 670.665 ]
+#    stds:
+#      S2L2A: [ 2106.761, 2141.107, 2038.973, 2134.138, 2085.321, 1889.926, 1820.257, 1871.918, 1753.829, 1797.379, 1434.261, 1334.311 ]
+#      S1RTC: [ 4.391, 4.459 ]
+#      DEM: [ 951.272 ]
+
+model:
+  class_path: terratorch.tasks.SemanticSegmentationTask
+  init_args:
+    model_factory: EncoderDecoderFactory
+    model_args:
+      backbone: terramind_v1_base
+      backbone_pretrained: true
+      backbone_modalities:
+        - S2L2A
+        - S1RTC
+        - DEM
+
+      # Apply temporal wrapper (docs: https://terrastackai.github.io/terratorch/stable/guide/temporal_wrapper/)
+      backbone_use_temporal: true
+      backbone_temporal_pooling: concat
+      backbone_temporal_n_timestamps: 4
+
+      necks:
+        - name: SelectIndices
+          indices: [2, 5, 8, 11]  # indices for terramind_v1_tiny, small, and base
+#          indices: [5, 11, 17, 23]  # large version
+        - name: ReshapeTokensToImage
+          remove_cls_token: False
+        - name: LearnedInterpolateToPyramidal
+
+      decoder: UNetDecoder
+#      decoder_channels: [256, 128, 64, 32] # tiny
+      decoder_channels: [512, 256, 128, 64] # base
+
+      head_dropout: 0.1
+      num_classes: 2
+    loss: dice
+    ignore_index: -1
+    freeze_backbone: false
+    freeze_decoder: false
+    class_weights: [0.342, 1.316]
+    # For prediction: overlap of 16 pixel on each side, 8 pixels dropped
+    tiled_inference_parameters:
+      crop: 256
+      stride: 208
+      batch_size: 64
+      delta: 8
+
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 1.e-4
+lr_scheduler:
+  class_path: ReduceLROnPlateau
+  init_args:
+    monitor: val/loss
+    factor: 0.5
+    patience: 2
+