(self, index)
| 195 | return max_num_frames, tensor_frms |
| 196 | |
| 197 | def __getitem__(self, index): |
| 198 | while True: |
| 199 | video_paths = self.video_list[index] |
| 200 | img_list = [ |
| 201 | os.path.join(self.data_root, filename) for filename in video_paths |
| 202 | ] |
| 203 | try: |
| 204 | num_frames, video_clip = self.read_img_list(img_list) |
| 205 | break |
| 206 | except Exception as e: |
| 207 | print("Broken data, skipping: {}".format(video_paths[-1])) |
| 208 | index = random.randint(0, self.length - 1) |
| 209 | continue |
| 210 | |
| 211 | # Add prefix |
| 212 | caption = self.captions_list[index] |
| 213 | prefix_prompt = self.prefix_prompt |
| 214 | if prefix_prompt != "": |
| 215 | prefix_prompt = prefix_prompt.strip() |
| 216 | prefix_prompt = prefix_prompt[0].upper() + prefix_prompt[1:] |
| 217 | if not prefix_prompt.endswith("."): |
| 218 | prefix_prompt += "." |
| 219 | |
| 220 | if self.p_drop_action_caption > 0 and random.random() < self.p_drop_action_caption: |
| 221 | caption = prefix_prompt |
| 222 | else: |
| 223 | caption = prefix_prompt + " " + caption |
| 224 | |
| 225 | # Traj |
| 226 | fut_traj = self.fut_traj_list[index] |
| 227 | fut_traj = torch.tensor(fut_traj, dtype=torch.float32) # [8, 3] |
| 228 | if self.p_mask_out_heading > 0 and random.random() < self.p_mask_out_heading: |
| 229 | fut_traj[:, -1] = 0 # mask out the heading |
| 230 | |
| 231 | if self.p_drop_traj > 0 and random.random() < self.p_drop_traj: # * drop the traj |
| 232 | fut_traj = torch.zeros_like(fut_traj) |
| 233 | |
| 234 | # Lidar pc token |
| 235 | lidar_pc_token = self.lidar_pc_token_list[index] |
| 236 | |
| 237 | # * For YouTube, human_drive_token can always apply |
| 238 | # * But for nuplan, it's only applied when fut_traj is dropped as zeros (free drive) |
| 239 | |
| 240 | with_human_drive_token = self.with_human_drive_token |
| 241 | if not self.always_apply_human_drive_token: |
| 242 | is_empty_traj = torch.all(fut_traj == 0).item() |
| 243 | with_human_drive_token = self.with_human_drive_token and is_empty_traj |
| 244 | |
| 245 | item = { |
| 246 | "with_traj": True, |
| 247 | "with_human_drive_token": with_human_drive_token, |
| 248 | "mp4": video_clip, |
| 249 | "txt": caption, |
| 250 | "num_frames": num_frames, |
| 251 | "fps": self.fps, # ? What's the use of fps? |
| 252 | "fut_traj": fut_traj, |
| 253 | "lidar_pc_token": lidar_pc_token |
| 254 | } |
nothing calls this directly
no test coverage detected