Remove reading frame lists from init function for live system

tcn_hpl/models/ptg_module.py

@@ -95,6 +95,7 @@ def __init__(
         actions_dict = dict()
         for a in actions:
             actions_dict[a.split()[1]] = int(a.split()[0])
+
         self.class_ids = list(actions_dict.values())
         self.classes = list(actions_dict.keys())
         self.action_id_to_str = dict(zip(self.class_ids, self.classes))
@@ -122,33 +123,8 @@ def __init__(
         self.validation_step_outputs_source_vid = []
         self.validation_step_outputs_source_frame = []
 
-        # Load val vidoes
-        vid_list_file_val = f"{self.hparams.data_dir}/splits/val.split1.bundle"
-        with open(vid_list_file_val, "r") as val_f:
-            self.val_videos = val_f.read().split("\n")[:-1]
-
-        self.val_frames = {}
-        for video in self.val_videos:
-            # Load frame filenames for the video
-            frame_list_file_val = f"{self.hparams.data_dir}/frames/{video}"
-            with open(frame_list_file_val, "r") as val_f:
-                val_fns = val_f.read().split("\n")[:-1]
-
-            self.val_frames[video[:-4]] = val_fns
-
-        # Load test vidoes
-        vid_list_file_tst = f"{self.hparams.data_dir}/splits/test.split1.bundle"
-        with open(vid_list_file_tst, "r") as test_f:
-            self.test_videos = test_f.read().split("\n")[:-1]
-
-        self.test_frames = {}
-        for video in self.test_videos:
-            # Load frame filenames for the video
-            frame_list_file_tst = f"{self.hparams.data_dir}/frames/{video}"
-            with open(frame_list_file_tst, "r") as test_f:
-                test_fns = test_f.read().split("\n")[:-1]
-
-            self.test_frames[video[:-4]] = test_fns
+        self.val_frames = None
+        self.test_frames = None
 
     def forward(self, x: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
         """Perform a forward pass through the model `self.net`.
@@ -288,6 +264,21 @@ def on_validation_epoch_end(self) -> None:
         all_source_vids = torch.concat(self.validation_step_outputs_source_vid)
         all_source_frames = torch.concat(self.validation_step_outputs_source_frame)
 
+        # Load val vidoes
+        if self.val_frames is None:
+            self.val_frames = {}
+            vid_list_file_val = f"{self.hparams.data_dir}/splits/val.split1.bundle"
+            with open(vid_list_file_val, "r") as val_f:
+                self.val_videos = val_f.read().split("\n")[:-1]
+
+            for video in self.val_videos:
+                # Load frame filenames for the video
+                frame_list_file_val = f"{self.hparams.data_dir}/frames/{video}"
+                with open(frame_list_file_val, "r") as val_f:
+                    val_fns = val_f.read().split("\n")[:-1]
+
+                self.val_frames[video[:-4]] = val_fns
+
         # Save results
         dset = kwcoco.CocoDataset()
         dset.fpath = f"{self.hparams.output_dir}/val_activity_preds_epoch{self.current_epoch}.mscoco.json"
@@ -373,6 +364,21 @@ def on_test_epoch_end(self) -> None:
         all_source_vids = torch.concat(self.validation_step_outputs_source_vid)
         all_source_frames = torch.concat(self.validation_step_outputs_source_frame)
 
+        # Load test vidoes
+        if self.test_frames is None:
+            self.test_frames = {}
+            vid_list_file_tst = f"{self.hparams.data_dir}/splits/test.split1.bundle"
+            with open(vid_list_file_tst, "r") as test_f:
+                self.test_videos = test_f.read().split("\n")[:-1]
+
+            for video in self.test_videos:
+                # Load frame filenames for the video
+                frame_list_file_tst = f"{self.hparams.data_dir}/frames/{video}"
+                with open(frame_list_file_tst, "r") as test_f:
+                    test_fns = test_f.read().split("\n")[:-1]
+
+                self.test_frames[video[:-4]] = test_fns
+
         # Save results
         dset = kwcoco.CocoDataset()
         dset.fpath = f"{self.hparams.output_dir}/test_activity_preds.mscoco.json"