(
self,
index: Optional[int] = None,
sequence_name: Optional[str] = None,
ids: Union[Iterable, None] = None,
)
| 219 | return self.get_data(index=index, ids=ids) |
| 220 | |
| 221 | def get_data( |
| 222 | self, |
| 223 | index: Optional[int] = None, |
| 224 | sequence_name: Optional[str] = None, |
| 225 | ids: Union[Iterable, None] = None, |
| 226 | ): |
| 227 | if sequence_name is None: |
| 228 | if index is None: |
| 229 | raise ValueError("Please specify either index or sequence_name") |
| 230 | sequence_name = self.sequence_list[index] |
| 231 | metadata = self.metadata[sequence_name] |
| 232 | category = self.category_map[sequence_name] |
| 233 | |
| 234 | if ids is None: |
| 235 | ids = np.arange(len(metadata)) |
| 236 | annos = [metadata[i] for i in ids] |
| 237 | if self.sort_by_filename: |
| 238 | annos = sorted(annos, key=lambda x: x["filepath"]) |
| 239 | |
| 240 | image_paths: list = [""] * len(annos) |
| 241 | focal_lengths: list = [0] * len(annos) |
| 242 | principal_points: list = [0] * len(annos) |
| 243 | |
| 244 | extrinsics: torch.Tensor = torch.eye(4, 4)[None].repeat(len(annos), 1, 1) |
| 245 | # intrinsics: torch.Tensor = torch.eye(3, 3)[None].repeat(len(annos), 1, 1) |
| 246 | |
| 247 | for idx, anno in enumerate(annos): |
| 248 | filepath = anno['filepath'] |
| 249 | impath = osp.join(self.CO3D_DIR, filepath) |
| 250 | |
| 251 | # focal_length = np.array(anno['focal_length']) |
| 252 | # principal_point = np.array(anno['principal_point']) |
| 253 | |
| 254 | extri_opencv = convert_pt3d_RT_to_opencv(anno["R"], anno["T"]) |
| 255 | |
| 256 | image_paths[idx] = impath |
| 257 | extrinsics[idx] = torch.tensor(extri_opencv) |
| 258 | focal_lengths[idx] = torch.tensor(anno['focal_length']) |
| 259 | principal_points[idx] = torch.tensor(anno['principal_point']) |
| 260 | |
| 261 | batch = {"seq_id": sequence_name, "category": category, "n": len(metadata), "ind": torch.tensor(ids)} |
| 262 | batch['image_paths'] = image_paths |
| 263 | |
| 264 | batch["extrs"] = extrinsics |
| 265 | batch["fl"] = torch.stack(focal_lengths) |
| 266 | batch["pp"] = torch.stack(principal_points) |
| 267 | |
| 268 | return batch |
| 269 | |
| 270 | def get_data_ori(self, index=None, sequence_name=None, ids=(0, 1), return_path = False): |
| 271 | if sequence_name is None: |
no test coverage detected