Remove 3 timestep restriction (#12)

- remove timestep restriction (67661a31b0739072008f148dcc51692e0733903a)
- modify readme for allowing different modes other than 3 timesteps (d397b2d16a5b2f6d6203c17adf7dd3237109ab9f)

Prithvi_run_inference.py

@@ -252,7 +252,7 @@ def main(data_files: List[str], yaml_file_path: str, checkpoint: str, output_dir
         params = yaml.safe_load(f)
 
     # data related
-    num_frames = params['num_frames']
+    num_frames = len(data_files)
     img_size = params['img_size']
     bands = params['bands']
     mean = params['data_mean']
@@ -272,8 +272,9 @@ def main(data_files: List[str], yaml_file_path: str, checkpoint: str, output_dir
 
     mask_ratio = params['mask_ratio'] if mask_ratio is None else mask_ratio
 
-    # We must have *num_frames* files to build one example!
-    assert len(data_files) == num_frames, "File list must be equal to expected number of frames."
+    print(f"\nTreating {len(data_files)} files as {len(data_files)} time steps from the same location\n")
+    if len(data_files) != 3:
+        print("The original model was trained for 3 time steps (expecting 3 files). \nResults with different numbers of timesteps may vary")
 
     if torch.cuda.is_available():
         device = torch.device('cuda')
@@ -310,7 +311,10 @@ def main(data_files: List[str], yaml_file_path: str, checkpoint: str, output_dir
     model.to(device)
 
     state_dict = torch.load(checkpoint, map_location=device)
-    model.load_state_dict(state_dict)
+    # discard fixed pos_embedding weight
+    del state_dict['pos_embed']
+    del state_dict['decoder_pos_embed']
+    model.load_state_dict(state_dict, strict=False)
     print(f"Loaded checkpoint from {checkpoint}")
 
     # Running model --------------------------------------------------------------------------------

README.md

@@ -33,7 +33,7 @@ The model follows the [original MAE repo](https://github.com/facebookresearch/ma
 4. adding infrared bands besides RGB
 
 ### Inference and demo
-There is an inference script (`Prithvi_run_inference.py`) that allows to run the image reconstruction on a set of three HLS images (see example below). These images have to be geotiff format, including the channels described above (Blue, Green, Red, Narrow NIR, SWIR 1, SWIR 2) in reflectance units. There is also a **demo** that leverages the same code [here](https://huggingface.co/spaces/ibm-nasa-geospatial/Prithvi-100M-demo).
+There is an inference script (`Prithvi_run_inference.py`) that allows to run the image reconstruction on a set of HLS images assumed to be from the same location at different time steps(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. There is also a **demo** that leverages the same code [here](https://huggingface.co/spaces/ibm-nasa-geospatial/Prithvi-100M-demo).
 
 ```
 python Prithvi_run_inference.py --data_files t1.tif t2.tif t3.tif --yaml_file_path /path/to/yaml/Prithvi_100.yaml --checkpoint /path/to/checkpoint/Prithvi_100.pth --output_dir /path/to/out/dir/ --mask_ratio 0.5