(self, directory_config, dataset_config, model_name, framerate=None, round_to_multiple=32, skip_dataset_validation=False)
| 446 | |
| 447 | class DirectoryDataset: |
| 448 | def __init__(self, directory_config, dataset_config, model_name, framerate=None, round_to_multiple=32, skip_dataset_validation=False): |
| 449 | self._set_defaults(directory_config, dataset_config) |
| 450 | self.directory_config = directory_config |
| 451 | self.dataset_config = dataset_config |
| 452 | if not skip_dataset_validation: |
| 453 | self.validate() |
| 454 | self.model_name = model_name |
| 455 | self.framerate = framerate |
| 456 | self.round_to_multiple = round_to_multiple |
| 457 | self.enable_ar_bucket = directory_config.get('enable_ar_bucket', dataset_config.get('enable_ar_bucket', False)) |
| 458 | # Configure directly from user-specified size buckets. |
| 459 | self.size_buckets = directory_config.get('size_buckets', dataset_config.get('size_buckets', None)) |
| 460 | self.use_size_buckets = (self.size_buckets is not None) |
| 461 | if self.use_size_buckets: |
| 462 | # sort size bucket from longest frame length to shortest |
| 463 | self.size_buckets.sort(key=lambda t: t[-1], reverse=True) |
| 464 | self.size_buckets = np.array(self.size_buckets) |
| 465 | self.size_bucket_datasets = [] |
| 466 | else: |
| 467 | self.resolutions = self._process_user_provided_resolutions( |
| 468 | directory_config.get('resolutions', dataset_config['resolutions']) |
| 469 | ) |
| 470 | self.resolutions = dedup_and_sort(self.resolutions) |
| 471 | self.ar_bucket_datasets = [] |
| 472 | self.shuffle = directory_config.get('cache_shuffle_num', dataset_config.get('cache_shuffle_num', 0)) |
| 473 | self.shuffle_metadata = directory_config['shuffle_metadata'] |
| 474 | self.directory_config['cache_shuffle_num'] = self.shuffle # Make accessible if it wasn't yet, for picking one out |
| 475 | self.shuffle_delimiter = directory_config.get('cache_shuffle_delimiter', dataset_config.get('cache_shuffle_delimiter', ", ")) |
| 476 | self.path = Path(self.directory_config['path']) |
| 477 | self.mask_path = Path(self.directory_config['mask_path']) if 'mask_path' in self.directory_config else None |
| 478 | self.control_path = Path(self.directory_config['control_path']) if 'control_path' in self.directory_config else None |
| 479 | # For testing. Default if a mask is missing. |
| 480 | self.default_mask_file = Path(self.directory_config['default_mask_file']) if 'default_mask_file' in self.directory_config else None |
| 481 | self.cache_dir = self.path / 'cache' / self.model_name |
| 482 | self.grouping_keys_json_file = self.cache_dir / 'metadata/grouping_keys.json' |
| 483 | self.skip_empty_caption = directory_config.get('skip_empty_caption', dataset_config.get('skip_empty_caption', True)) |
| 484 | |
| 485 | if not self.path.exists() or not self.path.is_dir(): |
| 486 | raise RuntimeError(f'Invalid path: {self.path}') |
| 487 | if self.mask_path is not None and (not self.mask_path.exists() or not self.mask_path.is_dir()): |
| 488 | raise RuntimeError(f'Invalid mask_path: {self.mask_path}') |
| 489 | if self.control_path is not None and (not self.control_path.exists() or not self.control_path.is_dir()): |
| 490 | raise RuntimeError(f'Invalid control_path: {self.control_path}') |
| 491 | if self.default_mask_file is not None and (not self.default_mask_file.exists() or not self.default_mask_file.is_file()): |
| 492 | raise RuntimeError(f'Invalid default_mask_file: {self.default_mask_file}') |
| 493 | |
| 494 | if self.use_size_buckets: |
| 495 | self.ars = np.array([w / h for w, h, _ in self.size_buckets]) |
| 496 | elif not self.enable_ar_bucket: |
| 497 | self.ars = np.array([1.0]) |
| 498 | elif ars := self.directory_config.get('ar_buckets', self.dataset_config.get('ar_buckets', None)): |
| 499 | self.ars = self._process_user_provided_ars(ars) |
| 500 | else: |
| 501 | min_ar = self.directory_config.get('min_ar', self.dataset_config['min_ar']) |
| 502 | max_ar = self.directory_config.get('max_ar', self.dataset_config['max_ar']) |
| 503 | num_ar_buckets = self.directory_config.get('num_ar_buckets', self.dataset_config['num_ar_buckets']) |
| 504 | self.ars = np.geomspace(min_ar, max_ar, num=num_ar_buckets) |
| 505 | self.ars = dedup_and_sort(self.ars) |
nothing calls this directly
no test coverage detected