| 24 | self._load_dataset() |
| 25 | |
| 26 | def _load_dataset(self): |
| 27 | split_dir = os.path.join(self.root, self.split) |
| 28 | |
| 29 | if not os.path.exists(split_dir): |
| 30 | raise FileNotFoundError(f"Dataset directory not found: {split_dir}") |
| 31 | |
| 32 | classes = sorted([d for d in os.listdir(split_dir) |
| 33 | if os.path.isdir(os.path.join(split_dir, d))]) |
| 34 | |
| 35 | self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(classes)} |
| 36 | |
| 37 | for class_name in classes: |
| 38 | class_dir = os.path.join(split_dir, class_name) |
| 39 | class_idx = self.class_to_idx[class_name] |
| 40 | |
| 41 | for img_name in os.listdir(class_dir): |
| 42 | if img_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')): |
| 43 | img_path = os.path.join(class_dir, img_name) |
| 44 | self.samples.append((img_path, class_idx)) |
| 45 | |
| 46 | def __len__(self) -> int: |
| 47 | return len(self.samples) |