| 25 | |
| 26 | |
| 27 | class YoloDataset(Dataset): |
| 28 | def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, phase: str = "train2017"): |
| 29 | augment_cfg = data_cfg.data_augment |
| 30 | self.image_size = data_cfg.image_size |
| 31 | phase_name = dataset_cfg.get(phase, phase) |
| 32 | self.batch_size = data_cfg.batch_size |
| 33 | self.dynamic_shape = getattr(data_cfg, "dynamic_shape", False) |
| 34 | self.base_size = mean(self.image_size) |
| 35 | |
| 36 | transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()] |
| 37 | self.transform = AugmentationComposer(transforms, self.image_size, self.base_size) |
| 38 | self.transform.get_more_data = self.get_more_data |
| 39 | self.img_paths, self.bboxes, self.ratios = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name)) |
| 40 | |
| 41 | def load_data(self, dataset_path: Path, phase_name: str): |
| 42 | """ |
| 43 | Loads data from a cache or generates a new cache for a specific dataset phase. |
| 44 | |
| 45 | Parameters: |
| 46 | dataset_path (Path): The root path to the dataset directory. |
| 47 | phase_name (str): The specific phase of the dataset (e.g., 'train', 'test') to load or generate data for. |
| 48 | |
| 49 | Returns: |
| 50 | dict: The loaded data from the cache for the specified phase. |
| 51 | """ |
| 52 | cache_path = dataset_path / f"{phase_name}.pache" |
| 53 | |
| 54 | if not cache_path.exists(): |
| 55 | logger.info(f":factory: Generating {phase_name} cache") |
| 56 | data = self.filter_data(dataset_path, phase_name, self.dynamic_shape) |
| 57 | torch.save(data, cache_path) |
| 58 | else: |
| 59 | try: |
| 60 | data = torch.load(cache_path, weights_only=False) |
| 61 | except Exception as e: |
| 62 | logger.error( |
| 63 | f":rotating_light: Failed to load the cache at '{cache_path}'.\n" |
| 64 | ":rotating_light: This may be caused by using cache from different other YOLO.\n" |
| 65 | ":rotating_light: Please clean the cache and try running again." |
| 66 | ) |
| 67 | raise e |
| 68 | logger.info(f":package: Loaded {phase_name} cache, there are {len(data)} data in total.") |
| 69 | return data |
| 70 | |
| 71 | def filter_data(self, dataset_path: Path, phase_name: str, sort_image: bool = False) -> list: |
| 72 | """ |
| 73 | Filters and collects dataset information by pairing images with their corresponding labels. |
| 74 | |
| 75 | Parameters: |
| 76 | images_path (Path): Path to the directory containing image files. |
| 77 | labels_path (str): Path to the directory containing label files. |
| 78 | sort_image (bool): If True, sorts the dataset by the width-to-height ratio of images in descending order. |
| 79 | |
| 80 | Returns: |
| 81 | list: A list of tuples, each containing the path to an image file and its associated segmentation as a tensor. |
| 82 | """ |
| 83 | images_path = dataset_path / "images" / phase_name |
| 84 | labels_path, data_type = locate_label_paths(dataset_path, phase_name) |