Update to use more parseable config

app.py

@@ -261,25 +261,18 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
         params = yaml.safe_load(f)
 
     # data related
-    num_frames = params['num_frames']
-    img_size = params['img_size']
-    bands = params['bands']
-    mean = params['data_mean']
-    std = params['data_std']
-
-    # model related
-    depth = params['depth']
-    patch_size = params['patch_size']
-    embed_dim = params['embed_dim']
-    num_heads = params['num_heads']
-    tubelet_size = params['tubelet_size']
-    decoder_embed_dim = params['decoder_embed_dim']
-    decoder_num_heads = params['decoder_num_heads']
-    decoder_depth = params['decoder_depth']
-
-    batch_size = params['batch_size']
-
-    mask_ratio = params['mask_ratio'] if mask_ratio is None else mask_ratio
+    train_params = params["train_params"]
+    num_frames = train_params['num_frames']
+    img_size = train_params['img_size']
+    bands = train_params['bands']
+    mean = train_params['data_mean']
+    std = train_params['data_std']
+
+    model_params = params["model_args"]
+
+    batch_size = 8
+
+    mask_ratio = train_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."
@@ -298,20 +291,7 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
     # Create model and load checkpoint -------------------------------------------------------------
 
     model = MaskedAutoencoderViT(
-            img_size=img_size,
-            patch_size=patch_size,
-            num_frames=num_frames,
-            tubelet_size=tubelet_size,
-            in_chans=len(bands),
-            embed_dim=embed_dim,
-            depth=depth,
-            num_heads=num_heads,
-            decoder_embed_dim=decoder_embed_dim,
-            decoder_depth=decoder_depth,
-            decoder_num_heads=decoder_num_heads,
-            mlp_ratio=4.,
-            norm_layer=functools.partial(torch.nn.LayerNorm, eps=1e-6),
-            norm_pix_loss=False)
+            **model_params)
 
     total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
     print(f"\n--> Model has {total_params:,} parameters.\n")