(
dataset_name: str,
dataset_info: T.Dict[str, T.Any],
input_camera_setting: T.Dict[str, T.Any] = None,
input_camera_trajectory_params: T.Dict[str, T.Any] = None,
output_camera_setting: T.Dict[str, T.Any] = None,
output_camera_trajectory_params: T.Dict[str, T.Any] = None,
rank: int = 0,
world_size: int = 1,
printout: bool = True,
imagenet_root_dir: str = 'datasets/imagenet',
)
| 430 | |
| 431 | |
| 432 | def get_dataset( |
| 433 | dataset_name: str, |
| 434 | dataset_info: T.Dict[str, T.Any], |
| 435 | input_camera_setting: T.Dict[str, T.Any] = None, |
| 436 | input_camera_trajectory_params: T.Dict[str, T.Any] = None, |
| 437 | output_camera_setting: T.Dict[str, T.Any] = None, |
| 438 | output_camera_trajectory_params: T.Dict[str, T.Any] = None, |
| 439 | rank: int = 0, |
| 440 | world_size: int = 1, |
| 441 | printout: bool = True, |
| 442 | imagenet_root_dir: str = 'datasets/imagenet', |
| 443 | ) -> T.Dict[str, T.Any]: |
| 444 | mesh_filename_dict = get_dataset_mesh_filename_config() |
| 445 | |
| 446 | if dataset_name not in mesh_filename_dict: |
| 447 | raise RuntimeError(f'{dataset_name} not in {get_dataset_mesh_filename()}') |
| 448 | |
| 449 | dataset_filename_dict = mesh_filename_dict[dataset_name] |
| 450 | |
| 451 | # download dataset |
| 452 | filename_dict = gather_and_clean_dataset( |
| 453 | dataset_name=dataset_name, |
| 454 | dataset_root_dir=dataset_filename_dict['dataset_root_dir'], |
| 455 | cleaned_root_dir=dataset_filename_dict.get('cleaned_root_dir', None), |
| 456 | clean_mesh=dataset_filename_dict.get('clean_mesh', False), |
| 457 | train_mesh_filenames=dataset_filename_dict['train'], |
| 458 | test_mesh_filenames=dataset_filename_dict.get('test', None), |
| 459 | rank=rank, |
| 460 | world_size=world_size, |
| 461 | printout=printout, |
| 462 | ) |
| 463 | train_mesh_filenames = filename_dict['train_mesh_filenames'] |
| 464 | test_mesh_filenames = filename_dict['test_mesh_filenames'] |
| 465 | |
| 466 | if world_size > 1: |
| 467 | torch.distributed.barrier() |
| 468 | |
| 469 | # make sure all files exist |
| 470 | for filename in train_mesh_filenames: |
| 471 | assert os.path.exists(filename), f'train: {filename} not exist' |
| 472 | if test_mesh_filenames is not None: |
| 473 | for filename in test_mesh_filenames: |
| 474 | assert os.path.exists(filename), f'test: {filename} not exist' |
| 475 | |
| 476 | # load mesh_filename to structures.Mesh |
| 477 | if printout: |
| 478 | print(f'Loading training meshes...', flush=True) |
| 479 | |
| 480 | if dataset_info.get('mix_meshes', False): |
| 481 | # combine multiple meshes |
| 482 | meshes: T.List[structures.Mesh] = load_and_mix_mesh_filename( |
| 483 | mesh_filenames=train_mesh_filenames, |
| 484 | mesh_scale=dataset_info.get('mesh_scale', 1.), |
| 485 | min_num_mesh=dataset_info.get('min_num_mesh', 1), |
| 486 | max_num_mesh=dataset_info.get('max_num_mesh', 2), |
| 487 | radius_scale=dataset_info.get('radius_scale', 2), |
| 488 | total_combined=dataset_info.get('total_combined', None), |
| 489 | printout=printout, |
nothing calls this directly
no test coverage detected