Update model training parameters and enhance dataset handling

commands.txt

@@ -47,8 +47,8 @@ venv/bin/python generate_keypoints.py \
 #3. Train (same as before, but make sure you use the venv python):
 venv/bin/python runner.py \
   --dataset include50 \
-  --use_augs \
   --model transformer \
+  --transformer_size large \
   --data_dir /home/ravijaanthony/Documents/dev/IIT/FYP/Code/INCLUDE/processed_data \
   --batch_size 8
 
@@ -57,9 +57,10 @@ venv/bin/python runner.py \
 venv/bin/python runner.py \
   --dataset include50 \
   --model transformer \
-  --transformer_size small \
+  --transformer_size large \
   --data_dir /home/ravijaanthony/Documents/dev/IIT/FYP/Code/INCLUDE/processed_data \
-  --use_pretrained evaluate
+  --batch_size 1
+  #--use_pretrained evaluate
 
 #4. A demo on a single video:
 venv/bin/python inference.py \

inference.py

@@ -90,8 +90,14 @@ def _pretrained_name(dataset: str, model_type: str, transformer_size: str) -> st
     return name
 
 
-def load_model(dataset: str, model_type: str, transformer_size: str, checkpoint_path: str | None):
-    label_map = load_label_map(dataset)
+def load_model(
+    dataset: str,
+    model_type: str,
+    transformer_size: str,
+    checkpoint_path: str | None,
+    label_map_path: str | None = None,
+):
+    label_map = load_json(label_map_path) if label_map_path else load_label_map(dataset)
     n_classes = len(label_map)
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

runner.py

@@ -42,7 +42,7 @@ parser.add_argument(
     help="location to save trained model",
 )
 parser.add_argument(
-    "--epochs", default=50, type=int, help="number of epochs to train the model"
+    "--epochs", default=150, type=int, help="number of epochs to train the model"
 )
 parser.add_argument("--batch_size", default=128, type=int, help="batch size of data")
 parser.add_argument(

utils.py

@@ -32,6 +32,8 @@ def get_experiment_name(args):
     if args.use_augs:
         exp_name += "augs_"
     exp_name += args.model
+    if args.model == "transformer":
+        exp_name += f"_{args.transformer_size}"
     return exp_name