Fix rename weight keys and fix inference

Prithvi-EO-V2-300M-TL-Sen1Floods11.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eb7db01b161d8919c7e61fed209d8633dabd49f5f81c785600790e43832344ba
-size 1276881720
+oid sha256:3675e9c2b52547de8ff8a19f4881c28573e6d4d2f0805d866f4fc48c1e517d60
+size 1276843350

config.yaml

@@ -101,13 +101,6 @@ data:
     - SWIR_1
     - SWIR_2
     train_transform:
-    - class_path: albumentations.Resize
-      init_args:
-        height: 448
-        width: 448
-        interpolation: 1
-        always_apply: false
-        p: 1
     - class_path: albumentations.RandomCrop
       init_args:
         height: 224
@@ -128,26 +121,12 @@ data:
         always_apply: true
         p: 1.0
     val_transform:
-    - class_path: albumentations.Resize
-      init_args:
-        height: 448
-        width: 448
-        interpolation: 1
-        always_apply: false
-        p: 1
     - class_path: albumentations.pytorch.ToTensorV2
       init_args:
         transpose_mask: false
         always_apply: true
         p: 1.0
     test_transform:
-    - class_path: albumentations.Resize
-      init_args:
-        height: 448
-        width: 448
-        interpolation: 1
-        always_apply: false
-        p: 1
     - class_path: albumentations.pytorch.ToTensorV2
       init_args:
         transpose_mask: false

inference.py

@@ -185,7 +185,7 @@ def run_model(input_data, temporal_coords, location_coords, model, datamodule, i
     for x in windows:
         # Apply standardization
         x = datamodule.test_transform(image=x.squeeze().numpy().transpose(1,2,0))
-        x = datamodule.aug(x['image'])
+        x = datamodule.aug(x)['image']
 
         with torch.no_grad():
             x = x.to(model.device)
@@ -237,7 +237,7 @@ def main(
     # Load model ---------------------------------------------------------------------------------
 
     lightning_model = LightningInferenceModel.from_config(config, checkpoint)
-    img_size = 512  # Size of Sen1Floods11
+    img_size = 512  # Size of Sen1Floods11 training
 
     # Loading data ---------------------------------------------------------------------------------
 

requirements.txt

@@ -3,4 +3,4 @@ torchvision
 timm
 einops
 rasterio
-git+https://github.com/IBM/terratorch.git
+terratorch==0.99.8