| 29 | |
| 30 | # * Repeat data here |
| 31 | def load_data_json(self, data_json, n_subset=None, ind_subset=None): |
| 32 | infos = load_json(data_json) |
| 33 | data_root = infos['meta']['data_root'] |
| 34 | clip_infos = infos['clips'] |
| 35 | if n_subset is not None and ind_subset is not None: |
| 36 | print("Using subset: {}/{}".format(ind_subset, n_subset)) |
| 37 | length_per_subset = math.ceil(len(clip_infos) / n_subset) |
| 38 | start_ind = ind_subset * length_per_subset |
| 39 | end_ind = (ind_subset + 1) * length_per_subset |
| 40 | clip_infos = clip_infos[start_ind:end_ind] |
| 41 | |
| 42 | for clip in tqdm(clip_infos): |
| 43 | |
| 44 | sample_path_tuple = (data_root, clip['folder_name'], clip['first_frame'], clip['end_frame']) |
| 45 | raw_caption = clip.get("flow_direction", "") # include static and highly static |
| 46 | |
| 47 | if self.use_psuedo_traj: |
| 48 | fut_traj = clip[self.traj_key][:self.n_fut_traj_points] |
| 49 | else: |
| 50 | fut_traj = torch.zeros((8, 3)) |
| 51 | |
| 52 | if self.exclude_highly_static and "Highly_Static" in raw_caption: |
| 53 | continue |
| 54 | |
| 55 | sample_caption = raw_caption |
| 56 | if self.merge_static and "Static" in sample_caption: |
| 57 | sample_caption = "Moving_Forward" # merging static and highly static to forward |
| 58 | |
| 59 | sample_caption = sample_caption.replace("_", " ").lower() # * MINOR FIX: Converting action to lowercase |
| 60 | sample_caption = sample_caption[0].upper() + sample_caption[1:] |
| 61 | if not sample_caption.endswith("."): |
| 62 | sample_caption += "." |
| 63 | |
| 64 | if self.n_repeat_of_actions is not None: |
| 65 | n_repeat = self.n_repeat_of_actions[raw_caption] |
| 66 | else: |
| 67 | n_repeat = 1 |
| 68 | |
| 69 | # * Repeat to achieve data weighting |
| 70 | sample_path_tuple = [sample_path_tuple] * n_repeat |
| 71 | sample_caption = [sample_caption] * n_repeat |
| 72 | sample_traj = [fut_traj] * n_repeat |
| 73 | |
| 74 | self.video_list.extend(sample_path_tuple) |
| 75 | self.captions_list.extend(sample_caption) |
| 76 | self.fut_traj_list.extend(sample_traj) |
| 77 | |
| 78 | def __getitem__(self, index): |
| 79 | while True: |