MCPcopy Create free account
hub / github.com/OpenDriveLab/ReSim / __getitem__

Method __getitem__

sat/data_share.py:197–255  ·  view source on GitHub ↗
(self, index)

Source from the content-addressed store, hash-verified

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 }

Callers

nothing calls this directly

Calls 2

read_img_listMethod · 0.95
printFunction · 0.50

Tested by

no test coverage detected