| 95 | |
| 96 | # * Repeat data here |
| 97 | def load_data_json(self, data_json, n_subset=None, ind_subset=None): |
| 98 | infos = load_json(data_json) |
| 99 | data_root = infos['meta']['data_root'] |
| 100 | self.data_root = data_root |
| 101 | |
| 102 | clip_infos = infos['clips'] |
| 103 | if self.token_json is not None: |
| 104 | token_json = load_json(self.token_json) |
| 105 | token_keep = token_json.keys() |
| 106 | clip_infos = [clip for clip in clip_infos if clip['lidar_pc_token'] in token_keep] |
| 107 | |
| 108 | if self.manual_total_length is not None: |
| 109 | if self.manual_total_length_mode == 'sequential': |
| 110 | clip_infos = clip_infos[:self.manual_total_length] |
| 111 | elif self.manual_total_length_mode == 'random': |
| 112 | clip_infos = random.sample(clip_infos, self.manual_total_length) |
| 113 | else: |
| 114 | raise ValueError("manual_total_length_mode should be either 'sequential' or 'random'") |
| 115 | |
| 116 | |
| 117 | if n_subset is not None and ind_subset is not None: |
| 118 | print("Using subset: {}/{}".format(ind_subset, n_subset)) |
| 119 | length_per_subset = math.ceil(len(clip_infos) / n_subset) |
| 120 | start_ind = ind_subset * length_per_subset |
| 121 | end_ind = (ind_subset + 1) * length_per_subset |
| 122 | clip_infos = clip_infos[start_ind:end_ind] |
| 123 | |
| 124 | |
| 125 | print("Using {} clips".format(len(clip_infos))) |
| 126 | |
| 127 | caption_key = "cmd" |
| 128 | |
| 129 | for clip in tqdm(clip_infos): |
| 130 | |
| 131 | if 'img_seq' in clip: |
| 132 | sample_seq = clip['img_seq'] |
| 133 | else: |
| 134 | sample_seq = clip['img_seq_his'] + clip['img_seq_fut'] |
| 135 | |
| 136 | raw_caption = clip.get(caption_key, "") # include static and highly static |
| 137 | if isinstance(raw_caption, int): |
| 138 | raw_caption = cmd_to_action[raw_caption] |
| 139 | |
| 140 | |
| 141 | fut_traj = clip[self.traj_key][:self.n_fut_traj_points] |
| 142 | |
| 143 | if self.p_use_extrapolate_traj > 0: |
| 144 | if 'extrapolated_traj_fut' in clip and random.random() < self.p_use_extrapolate_traj: |
| 145 | fut_traj = clip['extrapolated_traj_fut'][:self.n_fut_traj_points] |
| 146 | |
| 147 | |
| 148 | |
| 149 | token_key = 'lidar_pc_token' if 'lidar_pc_token' in clip else 'token' |
| 150 | lidar_pc_token = clip[token_key] |
| 151 | if not isinstance(lidar_pc_token, str): |
| 152 | lidar_pc_token = str(lidar_pc_token) |
| 153 | |
| 154 | sample_caption = raw_caption |