Switched from yaml to config.json

README.md

@@ -36,7 +36,7 @@ We provide a **demo** running Prithvi-EO-2.0-300M-TL [here](https://huggingface.
 There is also an inference script (`inference.py`) that allows to run the image reconstruction on a set of HLS images assumed to be from the same location at different timestamps (see example below). These should be provided in chronological order in geotiff format, including the channels described above (Blue, Green, Red, Narrow NIR, SWIR 1, SWIR 2) in reflectance units.
 
 ```
-python inference.py --data_files t1.tif t2.tif t3.tif t4.tif --output_dir output/ --input_indices <space separated 0-based indices of channels to select from input>
+python inference.py --data_files t1.tif t2.tif t3.tif t4.tif --input_indices <optional, space separated 0-based indices of the six Prithvi channels in your input>
 ```
 
 ## Finetuning

config.json

@@ -17,6 +17,7 @@
     "coords_scale_learn": true,
     "mask_ratio": 0.75,
     "norm_pix_loss": false,
+    "bands": ["B02", "B03", "B04", "B05", "B06", "B07"],
     "mean": [1087.0, 1342.0, 1433.0, 2734.0, 1958.0, 1363.0],
     "std": [2248.0, 2179.0, 2178.0, 1850.0, 1242.0, 1049.0],
     "origin_url": "https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL",

config.yaml

@@ -1,21 +0,0 @@
-DATA:
-  BANDS: [B02, B03, B04, B05, B06, B07]
-  INPUT_SIZE: [4, 224, 224]
-  MASK_RATIO: 0.75
-  MEAN: [1087.0, 1342.0, 1433.0, 2734.0, 1958.0, 1363.0]
-  STD: [2248.0, 2179.0, 2178.0, 1850.0, 1242.0, 1049.0]
-MODEL:
-  COORDS_DROP_RATE: 0.1
-  COORDS_ENCODING: [time, location]
-  COORDS_SCALE_LEARN: true
-  DECODER_DEPTH: 8
-  DECODER_EMBED_DIM: 512
-  DECODER_NUM_HEADS: 16
-  DEPTH: 24
-  DROP_CHANNELS_RATE: 0.0
-  EMBED_DIM: 1024
-  MLP_RATIO: 4.0
-  NAME: vit_l
-  NORM_PIX_LOSS: false
-  NUM_HEADS: 16
-  PATCH_SIZE: [1, 16, 16]

inference.py

@@ -304,18 +304,18 @@ def main(
 
     # Get parameters --------
 
+    import json
     with open(config_path, "r") as f:
-        config = yaml.safe_load(f)
+        config = yaml.safe_load(f)['pretrained_cfg']
 
     batch_size = 1
-    bands = config['DATA']['BANDS']
+    bands = config['bands']
     num_frames = len(data_files)
-    mean = config['DATA']['MEAN']
-    std = config['DATA']['STD']
-    coords_encoding = config['MODEL']['COORDS_ENCODING']
-    img_size = config['DATA']['INPUT_SIZE'][-1]
-
-    mask_ratio = mask_ratio or config['DATA']['MASK_RATIO']
+    mean = config['mean']
+    std = config['std']
+    coords_encoding = config['coords_encoding']
+    img_size = config['img_size']
+    mask_ratio = mask_ratio or config['mask_ratio']
 
     print(
         f"\nTreating {len(data_files)} files as {len(data_files)} time steps from the same location\n"
@@ -345,21 +345,13 @@ def main(
 
     # Create model and load checkpoint -------------------------------------------------------------
 
-    model = PrithviMAE(img_size=config['DATA']['INPUT_SIZE'][-2:],
-                       patch_size=config['MODEL']['PATCH_SIZE'],
-                       num_frames=num_frames,
-                       in_chans=len(bands),
-                       embed_dim=config['MODEL']['EMBED_DIM'],
-                       depth=config['MODEL']['DEPTH'],
-                       num_heads=config['MODEL']['NUM_HEADS'],
-                       decoder_embed_dim=config['MODEL']['DECODER_EMBED_DIM'],
-                       decoder_depth=config['MODEL']['DECODER_DEPTH'],
-                       decoder_num_heads=config['MODEL']['DECODER_NUM_HEADS'],
-                       mlp_ratio=config['MODEL']['MLP_RATIO'],
-                       norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
-                       norm_pix_loss=config['MODEL']['NORM_PIX_LOSS'],
-                       coords_encoding=coords_encoding,
-                       coords_scale_learn=config['MODEL']['COORDS_SCALE_LEARN'])
+    config.update(
+        coords_encoding=coords_encoding,
+        num_frames=num_frames,
+        in_chans=len(bands),
+    )
+
+    model = PrithviMAE(**config)
 
     total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
     print(f"\n--> Model has {total_params:,} parameters.\n")
@@ -487,11 +479,11 @@ if __name__ == "__main__":
         help="Path to the data files. Assumes multi-band files.",
     )
     parser.add_argument(
-        "--config_path",
+        "--config",
         "-c",
         type=str,
-        default="config.yaml",
-        help="Path to yaml file containing model training parameters.",
+        default="config.json",
+        help="Path to json file containing model training parameters.",
     )
     parser.add_argument(
         "--checkpoint",

prithvi_mae.py

@@ -240,7 +240,7 @@ class PrithviViT(nn.Module):
                  depth: int = 24,
                  num_heads: int = 16,
                  mlp_ratio: float = 4.,
-                 norm_layer: nn.Module = nn.LayerNorm,
+                 norm_layer: nn.Module = partial(torch.nn.LayerNorm, eps=1e-6),
                  coords_encoding: List[str] | None = None,
                  coords_scale_learn: bool = False,
                  encoder_only: bool = True,  # needed for timm
@@ -598,7 +598,7 @@ class PrithviMAE(nn.Module):
                  decoder_depth: int = 8,
                  decoder_num_heads: int = 16,
                  mlp_ratio: float = 4.,
-                 norm_layer: nn.Module = nn.LayerNorm,
+                 norm_layer: nn.Module = partial(torch.nn.LayerNorm, eps=1e-6),
                  norm_pix_loss: bool = False,
                  coords_encoding: List[str] | None = None,
                  coords_scale_learn: bool = False,