From b01296ea29afca71a10eff3507083268481c38da Mon Sep 17 00:00:00 2001 From: Hannah DeFazio Date: Thu, 26 Oct 2023 17:22:48 -0400 Subject: [PATCH] Remove reading frame lists from init function for live system --- tcn_hpl/models/ptg_module.py | 60 ++++++++++++++++++++---------------- 1 file changed, 33 insertions(+), 27 deletions(-) diff --git a/tcn_hpl/models/ptg_module.py b/tcn_hpl/models/ptg_module.py index 0bf58ccf6..0816ae6ae 100644 --- a/tcn_hpl/models/ptg_module.py +++ b/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"