dragon-ml-toolbox 12.13.0__py3-none-any.whl → 14.3.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of dragon-ml-toolbox might be problematic. Click here for more details.

Files changed (35) hide show
  1. {dragon_ml_toolbox-12.13.0.dist-info → dragon_ml_toolbox-14.3.0.dist-info}/METADATA +11 -2
  2. dragon_ml_toolbox-14.3.0.dist-info/RECORD +48 -0
  3. {dragon_ml_toolbox-12.13.0.dist-info → dragon_ml_toolbox-14.3.0.dist-info}/licenses/LICENSE-THIRD-PARTY.md +10 -0
  4. ml_tools/MICE_imputation.py +207 -5
  5. ml_tools/ML_callbacks.py +40 -8
  6. ml_tools/ML_datasetmaster.py +200 -261
  7. ml_tools/ML_evaluation.py +29 -17
  8. ml_tools/ML_evaluation_multi.py +13 -10
  9. ml_tools/ML_inference.py +14 -5
  10. ml_tools/ML_models.py +135 -55
  11. ml_tools/ML_models_advanced.py +323 -0
  12. ml_tools/ML_optimization.py +49 -36
  13. ml_tools/ML_trainer.py +560 -30
  14. ml_tools/ML_utilities.py +302 -4
  15. ml_tools/ML_vision_datasetmaster.py +1352 -0
  16. ml_tools/ML_vision_evaluation.py +260 -0
  17. ml_tools/ML_vision_inference.py +428 -0
  18. ml_tools/ML_vision_models.py +627 -0
  19. ml_tools/ML_vision_transformers.py +58 -0
  20. ml_tools/PSO_optimization.py +5 -1
  21. ml_tools/_ML_vision_recipe.py +88 -0
  22. ml_tools/__init__.py +1 -0
  23. ml_tools/_schema.py +96 -0
  24. ml_tools/custom_logger.py +37 -14
  25. ml_tools/data_exploration.py +576 -138
  26. ml_tools/keys.py +51 -1
  27. ml_tools/math_utilities.py +1 -1
  28. ml_tools/optimization_tools.py +65 -86
  29. ml_tools/serde.py +78 -17
  30. ml_tools/utilities.py +192 -3
  31. dragon_ml_toolbox-12.13.0.dist-info/RECORD +0 -41
  32. ml_tools/ML_simple_optimization.py +0 -413
  33. {dragon_ml_toolbox-12.13.0.dist-info → dragon_ml_toolbox-14.3.0.dist-info}/WHEEL +0 -0
  34. {dragon_ml_toolbox-12.13.0.dist-info → dragon_ml_toolbox-14.3.0.dist-info}/licenses/LICENSE +0 -0
  35. {dragon_ml_toolbox-12.13.0.dist-info → dragon_ml_toolbox-14.3.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,1352 @@
1
+ import torch
2
+ from torch.utils.data import Dataset, Subset
3
+ import numpy
4
+ from sklearn.model_selection import train_test_split
5
+ from typing import Union, Tuple, List, Optional, Callable, Dict, Any
6
+ from PIL import Image
7
+ from torchvision.datasets import ImageFolder
8
+ from torchvision import transforms
9
+ import torchvision.transforms.functional as TF
10
+ from pathlib import Path
11
+ import random
12
+ import json
13
+
14
+ from .ML_datasetmaster import _BaseMaker
15
+ from .path_manager import make_fullpath
16
+ from ._logger import _LOGGER
17
+ from ._script_info import _script_info
18
+ from .keys import VisionTransformRecipeKeys, ObjectDetectionKeys
19
+ from ._ML_vision_recipe import save_recipe
20
+ from .ML_vision_transformers import TRANSFORM_REGISTRY
21
+ from .custom_logger import custom_logger
22
+
23
+
24
+ __all__ = [
25
+ "VisionDatasetMaker",
26
+ "SegmentationDatasetMaker",
27
+ "ObjectDetectionDatasetMaker"
28
+ ]
29
+
30
+
31
+ # --- VisionDatasetMaker ---
32
+ class VisionDatasetMaker(_BaseMaker):
33
+ """
34
+ Creates processed PyTorch datasets for computer vision tasks from an
35
+ image folder directory.
36
+
37
+ Supports two modes:
38
+ 1. `from_folder()`: Loads from one directory and splits into train/val/test.
39
+ 2. `from_folders()`: Loads from pre-split train/val/test directories.
40
+
41
+ Uses online augmentations per epoch (image augmentation without creating new files).
42
+ """
43
+ def __init__(self):
44
+ """
45
+ Typically not called directly. Use the class methods `from_folder()` or `from_folders()` to create an instance.
46
+ """
47
+ super().__init__()
48
+ self._full_dataset: Optional[ImageFolder] = None
49
+ self.labels: Optional[List[int]] = None
50
+ self.class_map: Optional[dict[str,int]] = None
51
+
52
+ self._is_split = False
53
+ self._are_transforms_configured = False
54
+ self._val_recipe_components = None
55
+ self._has_mean_std: bool = False
56
+
57
+ @classmethod
58
+ def from_folder(cls, root_dir: Union[str,Path]) -> 'VisionDatasetMaker':
59
+ """
60
+ Creates a maker instance from a single root directory of images.
61
+
62
+ This method assumes a single directory (e.g., 'data/') that
63
+ contains class subfolders (e.g., 'data/cat/', 'data/dog/').
64
+
65
+ The dataset will be loaded in its entirety, and you MUST call
66
+ `.split_data()` afterward to create train/validation/test sets.
67
+
68
+ Args:
69
+ root_dir (str | Path): The path to the root directory containing class subfolders.
70
+
71
+ Returns:
72
+ VisionDatasetMaker: A new instance with the full dataset loaded.
73
+ """
74
+ root_path = make_fullpath(root_dir, enforce="directory")
75
+ # Load with NO transform. We get PIL Images.
76
+ full_dataset = ImageFolder(root=root_path, transform=None)
77
+ _LOGGER.info(f"Found {len(full_dataset)} images in {len(full_dataset.classes)} classes.")
78
+
79
+ maker = cls()
80
+ maker._full_dataset = full_dataset
81
+ maker.labels = [s[1] for s in full_dataset.samples]
82
+ maker.class_map = full_dataset.class_to_idx
83
+ return maker
84
+
85
+ @classmethod
86
+ def from_folders(cls,
87
+ train_dir: Union[str,Path],
88
+ val_dir: Union[str,Path],
89
+ test_dir: Optional[Union[str,Path]] = None) -> 'VisionDatasetMaker':
90
+ """
91
+ Creates a maker instance from separate, pre-split directories.
92
+
93
+ This method is used when you already have 'train', 'val', and
94
+ optionally 'test' folders, each containing class subfolders.
95
+ It bypasses the need for `.split_data()`.
96
+
97
+ Args:
98
+ train_dir (str | Path): Path to the training data directory.
99
+ val_dir (str | Path): Path to the validation data directory.
100
+ test_dir (str | Path | None): Path to the test data directory.
101
+
102
+ Returns:
103
+ VisionDatasetMaker: A new, pre-split instance.
104
+
105
+ Raises:
106
+ ValueError: If the classes found in train, val, or test directories are inconsistent.
107
+ """
108
+ train_path = make_fullpath(train_dir, enforce="directory")
109
+ val_path = make_fullpath(val_dir, enforce="directory")
110
+
111
+ _LOGGER.info("Loading data from separate directories.")
112
+ # Load with NO transform. We get PIL Images.
113
+ train_ds = ImageFolder(root=train_path, transform=None)
114
+ val_ds = ImageFolder(root=val_path, transform=None)
115
+
116
+ # Check for class consistency
117
+ if train_ds.class_to_idx != val_ds.class_to_idx:
118
+ _LOGGER.error("Train and validation directories have different or inconsistent classes.")
119
+ raise ValueError()
120
+
121
+ maker = cls()
122
+ maker._train_dataset = train_ds
123
+ maker._val_dataset = val_ds
124
+ maker.class_map = train_ds.class_to_idx
125
+
126
+ if test_dir:
127
+ test_path = make_fullpath(test_dir, enforce="directory")
128
+ test_ds = ImageFolder(root=test_path, transform=None)
129
+ if train_ds.class_to_idx != test_ds.class_to_idx:
130
+ _LOGGER.error("Train and test directories have different or inconsistent classes.")
131
+ raise ValueError()
132
+ maker._test_dataset = test_ds
133
+ _LOGGER.info(f"Loaded: {len(train_ds)} train, {len(val_ds)} val, {len(test_ds)} test images.")
134
+ else:
135
+ _LOGGER.info(f"Loaded: {len(train_ds)} train, {len(val_ds)} val images.")
136
+
137
+ maker._is_split = True # Mark as "split" since data is pre-split
138
+ return maker
139
+
140
+ @staticmethod
141
+ def inspect_folder(path: Union[str, Path]):
142
+ """
143
+ Logs a report of the types, sizes, and channels of image files
144
+ found in the directory and its subdirectories.
145
+
146
+ This is a utility method to help diagnose potential dataset
147
+ issues (e.g., mixed image modes, corrupted files) before loading.
148
+
149
+ Args:
150
+ path (str, Path): The directory path to inspect.
151
+ """
152
+ path_obj = make_fullpath(path)
153
+
154
+ non_image_files = set()
155
+ img_types = set()
156
+ img_sizes = set()
157
+ img_channels = set()
158
+ img_counter = 0
159
+
160
+ _LOGGER.info(f"Inspecting folder: {path_obj}...")
161
+ # Use rglob to recursively find all files
162
+ for filepath in path_obj.rglob('*'):
163
+ if filepath.is_file():
164
+ try:
165
+ # Using PIL to open is a more reliable check
166
+ with Image.open(filepath) as img:
167
+ img_types.add(img.format)
168
+ img_sizes.add(img.size)
169
+ img_channels.update(img.getbands())
170
+ img_counter += 1
171
+ except (IOError, SyntaxError):
172
+ non_image_files.add(filepath.name)
173
+
174
+ if non_image_files:
175
+ _LOGGER.warning(f"Non-image or corrupted files found and ignored: {non_image_files}")
176
+
177
+ report = (
178
+ f"\n--- Inspection Report for '{path_obj.name}' ---\n"
179
+ f"Total images found: {img_counter}\n"
180
+ f"Image formats: {img_types or 'None'}\n"
181
+ f"Image sizes (WxH): {img_sizes or 'None'}\n"
182
+ f"Image channels (bands): {img_channels or 'None'}\n"
183
+ f"--------------------------------------"
184
+ )
185
+ print(report)
186
+
187
+ def split_data(self, val_size: float = 0.2, test_size: float = 0.0,
188
+ stratify: bool = True, random_state: Optional[int] = None) -> 'VisionDatasetMaker':
189
+ """
190
+ Splits the dataset into train, validation, and optional test sets.
191
+
192
+ This method MUST be called if you used `from_folder()`. It has no effect if you used `from_folders()`.
193
+
194
+ Args:
195
+ val_size (float): Proportion of the dataset to reserve for
196
+ validation (e.g., 0.2 for 20%).
197
+ test_size (float): Proportion of the dataset to reserve for
198
+ testing.
199
+ stratify (bool): If True, splits are performed in a stratified
200
+ fashion, preserving the class distribution.
201
+ random_state (int | None): Seed for the random number generator for reproducible splits.
202
+
203
+ Returns:
204
+ VisionDatasetMaker: The same instance, now with datasets split.
205
+
206
+ Raises:
207
+ ValueError: If `val_size` and `test_size` sum to 1.0 or more.
208
+ """
209
+ if self._is_split:
210
+ _LOGGER.warning("Data has already been split.")
211
+ return self
212
+
213
+ if val_size + test_size >= 1.0:
214
+ _LOGGER.error("The sum of val_size and test_size must be less than 1.")
215
+ raise ValueError()
216
+
217
+ if not self._full_dataset:
218
+ _LOGGER.error("There is no dataset to split.")
219
+ raise ValueError()
220
+
221
+ indices = list(range(len(self._full_dataset)))
222
+ labels_for_split = self.labels if stratify else None
223
+
224
+ train_indices, val_test_indices = train_test_split(
225
+ indices, test_size=(val_size + test_size), random_state=random_state, stratify=labels_for_split
226
+ )
227
+
228
+ if not self.labels:
229
+ _LOGGER.error("Error when getting full dataset labels.")
230
+ raise ValueError()
231
+
232
+ if test_size > 0:
233
+ val_test_labels = [self.labels[i] for i in val_test_indices]
234
+ stratify_val_test = val_test_labels if stratify else None
235
+ val_indices, test_indices = train_test_split(
236
+ val_test_indices, test_size=(test_size / (val_size + test_size)),
237
+ random_state=random_state, stratify=stratify_val_test
238
+ )
239
+ self._test_dataset = Subset(self._full_dataset, test_indices)
240
+ _LOGGER.info(f"Test set created with {len(self._test_dataset)} images.")
241
+ else:
242
+ val_indices = val_test_indices
243
+
244
+ self._train_dataset = Subset(self._full_dataset, train_indices)
245
+ self._val_dataset = Subset(self._full_dataset, val_indices)
246
+ self._is_split = True
247
+
248
+ _LOGGER.info(f"Data split into: \n- Training: {len(self._train_dataset)} images \n- Validation: {len(self._val_dataset)} images")
249
+ return self
250
+
251
+ def configure_transforms(self, resize_size: int = 256, crop_size: int = 224,
252
+ mean: Optional[List[float]] = [0.485, 0.456, 0.406],
253
+ std: Optional[List[float]] = [0.229, 0.224, 0.225],
254
+ pre_transforms: Optional[List[Callable]] = None,
255
+ extra_train_transforms: Optional[List[Callable]] = None) -> 'VisionDatasetMaker':
256
+ """
257
+ Configures and applies the image transformations and augmentations.
258
+
259
+ This method must be called AFTER data is loaded and split.
260
+
261
+ It sets up two pipelines:
262
+ 1. **Training Pipeline:** Includes random augmentations like
263
+ `RandomResizedCrop` and `RandomHorizontalFlip` (plus any
264
+ `extra_train_transforms`) for online augmentation.
265
+ 2. **Validation/Test Pipeline:** A deterministic pipeline using
266
+ `Resize` and `CenterCrop` for consistent evaluation.
267
+
268
+ Both pipelines finish with `ToTensor` and `Normalize`.
269
+
270
+ Args:
271
+ resize_size (int): The size to resize the smallest edge to
272
+ for validation/testing.
273
+ crop_size (int): The target size (square) for the final
274
+ cropped image.
275
+ mean (List[float]): The mean values for normalization (e.g., ImageNet mean).
276
+ std (List[float]): The standard deviation values for normalization (e.g., ImageNet std).
277
+ extra_train_transforms (List[Callable] | None): A list of additional torchvision transforms to add to the end of the training transformations.
278
+ pre_transforms (List[Callable] | None): An list of transforms to be applied at the very beginning of the transformations for all sets.
279
+
280
+ Returns:
281
+ VisionDatasetMaker: The same instance, with transforms applied.
282
+
283
+ Raises:
284
+ RuntimeError: If called before data is split.
285
+ """
286
+ if not self._is_split:
287
+ _LOGGER.error("Transforms must be configured AFTER splitting data (or using `from_folders`). Call .split_data() first if using `from_folder`.")
288
+ raise RuntimeError()
289
+
290
+ if (mean is None and std is not None) or (mean is not None and std is None):
291
+ _LOGGER.error(f"'mean' and 'std' must be both None or both defined, but only one was provided.")
292
+ raise ValueError()
293
+
294
+ # --- Define Transform Pipelines ---
295
+ # These now MUST include ToTensor and Normalize, as the ImageFolder was loaded with transform=None.
296
+
297
+ # --- Store components for validation recipe ---
298
+ self._val_recipe_components = {
299
+ VisionTransformRecipeKeys.PRE_TRANSFORMS: pre_transforms or [],
300
+ VisionTransformRecipeKeys.RESIZE_SIZE: resize_size,
301
+ VisionTransformRecipeKeys.CROP_SIZE: crop_size,
302
+ }
303
+
304
+ if mean is not None and std is not None:
305
+ self._val_recipe_components.update({
306
+ VisionTransformRecipeKeys.MEAN: mean,
307
+ VisionTransformRecipeKeys.STD: std
308
+ })
309
+ self._has_mean_std = True
310
+
311
+ base_pipeline = []
312
+ if pre_transforms:
313
+ base_pipeline.extend(pre_transforms)
314
+
315
+ # Base augmentations for training
316
+ base_train_transforms = [
317
+ transforms.RandomResizedCrop(crop_size),
318
+ transforms.RandomHorizontalFlip()
319
+ ]
320
+ if extra_train_transforms:
321
+ base_train_transforms.extend(extra_train_transforms)
322
+
323
+ # Final conversion and normalization
324
+ final_transforms: list[Callable] = [
325
+ transforms.ToTensor()
326
+ ]
327
+
328
+ if self._has_mean_std:
329
+ final_transforms.append(transforms.Normalize(mean=mean, std=std))
330
+
331
+ # Build the val/test pipeline
332
+ val_transform_list = [
333
+ *base_pipeline, # Apply pre_transforms first
334
+ transforms.Resize(resize_size),
335
+ transforms.CenterCrop(crop_size),
336
+ *final_transforms
337
+ ]
338
+
339
+ # Build the train pipeline
340
+ train_transform_list = [
341
+ *base_pipeline, # Apply pre_transforms first
342
+ *base_train_transforms,
343
+ *final_transforms
344
+ ]
345
+
346
+ val_transform = transforms.Compose(val_transform_list)
347
+ train_transform = transforms.Compose(train_transform_list)
348
+
349
+ # --- Apply Transforms using the Wrapper ---
350
+ # This correctly assigns the transform regardless of whether the dataset is a Subset (from_folder) or an ImageFolder (from_folders).
351
+
352
+ self._train_dataset = _DatasetTransformer(self._train_dataset, train_transform) # type: ignore
353
+ self._val_dataset = _DatasetTransformer(self._val_dataset, val_transform) # type: ignore
354
+ if self._test_dataset:
355
+ self._test_dataset = _DatasetTransformer(self._test_dataset, val_transform) # type: ignore
356
+
357
+ self._are_transforms_configured = True
358
+ _LOGGER.info("Image transforms configured and applied.")
359
+ return self
360
+
361
+ def get_datasets(self) -> Tuple[Dataset, ...]:
362
+ """
363
+ Returns the final train, validation, and optional test datasets.
364
+
365
+ This is the final step, used to retrieve the datasets for use in
366
+ a `MLTrainer` or `DataLoader`.
367
+
368
+ Returns:
369
+ (Tuple[Dataset, ...]): A tuple containing the (train, val)
370
+ or (train, val, test) datasets.
371
+
372
+ Raises:
373
+ RuntimeError: If called before data is split.
374
+ UserWarning: If called before transforms are configured.
375
+ """
376
+ if not self._is_split:
377
+ _LOGGER.error("Data has not been split. Call .split_data() first.")
378
+ raise RuntimeError()
379
+ if not self._are_transforms_configured:
380
+ _LOGGER.warning("Transforms have not been configured.")
381
+
382
+ if self._test_dataset:
383
+ return self._train_dataset, self._val_dataset, self._test_dataset
384
+ return self._train_dataset, self._val_dataset
385
+
386
+ def save_transform_recipe(self, filepath: Union[str, Path]) -> None:
387
+ """
388
+ Saves the validation transform pipeline as a JSON recipe file.
389
+
390
+ This recipe can be loaded by the PyTorchVisionInferenceHandler
391
+ to ensure identical preprocessing.
392
+
393
+ Args:
394
+ filepath (str | Path): The path to save the .json recipe file.
395
+ """
396
+ if not self._are_transforms_configured:
397
+ _LOGGER.error("Transforms are not configured. Call .configure_transforms() first.")
398
+ raise RuntimeError()
399
+
400
+ recipe: Dict[str, Any] = {
401
+ VisionTransformRecipeKeys.TASK: "classification",
402
+ VisionTransformRecipeKeys.PIPELINE: []
403
+ }
404
+
405
+ components = self._val_recipe_components
406
+
407
+ if not components:
408
+ _LOGGER.error(f"Error getting the transformers recipe for validation set.")
409
+ raise ValueError()
410
+
411
+ # validate path
412
+ file_path = make_fullpath(filepath, make=True, enforce="file")
413
+
414
+ # 1. Handle pre_transforms
415
+ for t in components[VisionTransformRecipeKeys.PRE_TRANSFORMS]:
416
+ t_name = t.__class__.__name__
417
+ if t_name in TRANSFORM_REGISTRY:
418
+ recipe[VisionTransformRecipeKeys.PIPELINE].append({
419
+ VisionTransformRecipeKeys.NAME: t_name,
420
+ VisionTransformRecipeKeys.KWARGS: getattr(t, VisionTransformRecipeKeys.KWARGS, {})
421
+ })
422
+ else:
423
+ _LOGGER.warning(f"Skipping unknown pre_transform '{t_name}' in recipe. Not in TRANSFORM_REGISTRY.")
424
+
425
+ # 2. Add standard transforms
426
+ recipe[VisionTransformRecipeKeys.PIPELINE].extend([
427
+ {VisionTransformRecipeKeys.NAME: "Resize", "kwargs": {"size": components[VisionTransformRecipeKeys.RESIZE_SIZE]}},
428
+ {VisionTransformRecipeKeys.NAME: "CenterCrop", "kwargs": {"size": components[VisionTransformRecipeKeys.CROP_SIZE]}},
429
+ {VisionTransformRecipeKeys.NAME: "ToTensor", "kwargs": {}}
430
+ ])
431
+
432
+ if self._has_mean_std:
433
+ recipe[VisionTransformRecipeKeys.PIPELINE].append(
434
+ {VisionTransformRecipeKeys.NAME: "Normalize", "kwargs": {
435
+ "mean": components[VisionTransformRecipeKeys.MEAN],
436
+ "std": components[VisionTransformRecipeKeys.STD]
437
+ }}
438
+ )
439
+
440
+ # 3. Save the file
441
+ save_recipe(recipe, file_path)
442
+
443
+ def save_class_map(self, save_dir: Union[str,Path]) -> dict[str,int]:
444
+ """
445
+ Saves the class to index mapping {str: int} to a directory.
446
+ """
447
+ if not self.class_map:
448
+ _LOGGER.error(f"Class to index mapping is empty.")
449
+ raise ValueError()
450
+
451
+ custom_logger(data=self.class_map,
452
+ save_directory=save_dir,
453
+ log_name="Class_to_Index",
454
+ add_timestamp=False,
455
+ dict_as="json")
456
+
457
+ return self.class_map
458
+
459
+
460
+ class _DatasetTransformer(Dataset):
461
+ """
462
+ Internal wrapper class to apply a specific transform pipeline to any
463
+ dataset (e.g., a full ImageFolder or a Subset).
464
+ """
465
+ def __init__(self, dataset: Dataset, transform: Optional[transforms.Compose] = None):
466
+ self.dataset = dataset
467
+ self.transform = transform
468
+
469
+ # --- Propagate attributes for inspection ---
470
+ # For ImageFolder
471
+ if hasattr(dataset, 'class_to_idx'):
472
+ self.class_to_idx = getattr(dataset, 'class_to_idx')
473
+ if hasattr(dataset, 'classes'):
474
+ self.classes = getattr(dataset, 'classes')
475
+ # For Subset
476
+ if hasattr(dataset, 'indices'):
477
+ self.indices = getattr(dataset, 'indices')
478
+ if hasattr(dataset, 'dataset'):
479
+ # This allows access to the *original* full dataset
480
+ self.original_dataset = getattr(dataset, 'dataset')
481
+
482
+ def __getitem__(self, index):
483
+ # Get the original data (e.g., PIL Image, label)
484
+ x, y = self.dataset[index]
485
+
486
+ # Apply the specific transform for this dataset
487
+ if self.transform:
488
+ x = self.transform(x)
489
+ return x, y
490
+
491
+ def __len__(self):
492
+ return len(self.dataset) # type: ignore
493
+
494
+
495
+ # --- Segmentation dataset ----
496
+ class _SegmentationDataset(Dataset):
497
+ """
498
+ Internal helper class to load image-mask pairs.
499
+
500
+ Loads images as RGB and masks as 'L' (grayscale, 8-bit integer pixels).
501
+ """
502
+ def __init__(self, image_paths: List[Path], mask_paths: List[Path], transform: Optional[Callable] = None):
503
+ self.image_paths = image_paths
504
+ self.mask_paths = mask_paths
505
+ self.transform = transform
506
+
507
+ # --- Propagate 'classes' if they exist (for MLTrainer) ---
508
+ self.classes: List[str] = []
509
+
510
+ def __len__(self):
511
+ return len(self.image_paths)
512
+
513
+ def __getitem__(self, idx):
514
+ img_path = self.image_paths[idx]
515
+ mask_path = self.mask_paths[idx]
516
+
517
+ try:
518
+ # Open as PIL Images. Masks should be 'L'
519
+ image = Image.open(img_path).convert("RGB")
520
+ mask = Image.open(mask_path).convert("L")
521
+ except Exception as e:
522
+ _LOGGER.error(f"Error loading sample #{idx}: {img_path.name} / {mask_path.name}. Error: {e}")
523
+ # Return empty tensors
524
+ return torch.empty(3, 224, 224), torch.empty(224, 224, dtype=torch.long)
525
+
526
+ if self.transform:
527
+ image, mask = self.transform(image, mask)
528
+
529
+ return image, mask
530
+
531
+
532
+ # Internal Paired Transform Helpers
533
+ class _PairedCompose:
534
+ """A 'Compose' for paired image/mask transforms."""
535
+ def __init__(self, transforms: List[Callable]):
536
+ self.transforms = transforms
537
+
538
+ def __call__(self, image: Any, mask: Any) -> Tuple[Any, Any]:
539
+ for t in self.transforms:
540
+ image, mask = t(image, mask)
541
+ return image, mask
542
+
543
+ class _PairedToTensor:
544
+ """Converts a PIL Image pair (image, mask) to Tensors."""
545
+ def __call__(self, image: Image.Image, mask: Image.Image) -> Tuple[torch.Tensor, torch.Tensor]:
546
+ # Use new variable names to satisfy the linter
547
+ image_tensor = TF.to_tensor(image)
548
+ # Convert mask to LongTensor, not float.
549
+ # This creates a [H, W] tensor of integer class IDs.
550
+ mask_tensor = torch.from_numpy(numpy.array(mask, dtype=numpy.int64))
551
+ return image_tensor, mask_tensor
552
+
553
+ class _PairedNormalize:
554
+ """Normalizes the image tensor and leaves the mask untouched."""
555
+ def __init__(self, mean: List[float], std: List[float]):
556
+ self.normalize = transforms.Normalize(mean, std)
557
+
558
+ def __call__(self, image: torch.Tensor, mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
559
+ image = self.normalize(image)
560
+ return image, mask
561
+
562
+ class _PairedResize:
563
+ """Resizes an image and mask to the same size."""
564
+ def __init__(self, size: int):
565
+ self.size = [size, size]
566
+
567
+ def __call__(self, image: Image.Image, mask: Image.Image) -> Tuple[Image.Image, Image.Image]:
568
+ # Use new variable names to avoid linter confusion
569
+ resized_image = TF.resize(image, self.size, interpolation=TF.InterpolationMode.BILINEAR) # type: ignore
570
+ # Use NEAREST for mask to avoid interpolating class IDs (e.g., 1.5)
571
+ resized_mask = TF.resize(mask, self.size, interpolation=TF.InterpolationMode.NEAREST) # type: ignore
572
+ return resized_image, resized_mask # type: ignore
573
+
574
+ class _PairedCenterCrop:
575
+ """Center-crops an image and mask to the same size."""
576
+ def __init__(self, size: int):
577
+ self.size = [size, size]
578
+
579
+ def __call__(self, image: Image.Image, mask: Image.Image) -> Tuple[Image.Image, Image.Image]:
580
+ cropped_image = TF.center_crop(image, self.size) # type: ignore
581
+ cropped_mask = TF.center_crop(mask, self.size) # type: ignore
582
+ return cropped_image, cropped_mask # type: ignore
583
+
584
+ class _PairedRandomHorizontalFlip:
585
+ """Applies the same random horizontal flip to both image and mask."""
586
+ def __init__(self, p: float = 0.5):
587
+ self.p = p
588
+
589
+ def __call__(self, image: Image.Image, mask: Image.Image) -> Tuple[Image.Image, Image.Image]:
590
+ if random.random() < self.p:
591
+ flipped_image = TF.hflip(image) # type: ignore
592
+ flipped_mask = TF.hflip(mask) # type: ignore
593
+ return flipped_image, flipped_mask # type: ignore
594
+
595
+ class _PairedRandomResizedCrop:
596
+ """Applies the same random resized crop to both image and mask."""
597
+ def __init__(self, size: int, scale: Tuple[float, float]=(0.08, 1.0), ratio: Tuple[float, float]=(3./4., 4./3.)):
598
+ self.size = [size, size]
599
+ self.scale = scale
600
+ self.ratio = ratio
601
+ self.interpolation = TF.InterpolationMode.BILINEAR
602
+ self.mask_interpolation = TF.InterpolationMode.NEAREST
603
+
604
+ def __call__(self, image: Image.Image, mask: Image.Image) -> Tuple[Image.Image, Image.Image]:
605
+ # Get parameters for the random crop
606
+ # Convert scale/ratio tuples to lists to satisfy the linter's type hint
607
+ i, j, h, w = transforms.RandomResizedCrop.get_params(image, list(self.scale), list(self.ratio)) # type: ignore
608
+
609
+ # Apply the crop with the SAME parameters and use new variable names
610
+ cropped_image = TF.resized_crop(image, i, j, h, w, self.size, self.interpolation) # type: ignore
611
+ cropped_mask = TF.resized_crop(mask, i, j, h, w, self.size, self.mask_interpolation) # type: ignore
612
+
613
+ return cropped_image, cropped_mask # type: ignore
614
+
615
+ # --- SegmentationDatasetMaker ---
616
+ class SegmentationDatasetMaker(_BaseMaker):
617
+ """
618
+ Creates processed PyTorch datasets for segmentation from image and mask folders.
619
+
620
+ This maker finds all matching image-mask pairs from two directories,
621
+ splits them, and applies identical transformations (including augmentations)
622
+ to both the image and its corresponding mask.
623
+
624
+ Workflow:
625
+ 1. `maker = SegmentationDatasetMaker.from_folders(img_dir, mask_dir)`
626
+ 2. `maker.set_class_map({'background': 0, 'road': 1})`
627
+ 3. `maker.split_data(val_size=0.2)`
628
+ 4. `maker.configure_transforms(crop_size=256)`
629
+ 5. `train_ds, val_ds = maker.get_datasets()`
630
+ """
631
+ IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff')
632
+
633
+ def __init__(self):
634
+ """
635
+ Typically not called directly. Use the class method `from_folders()` to create an instance.
636
+ """
637
+ super().__init__()
638
+ self.image_paths: List[Path] = []
639
+ self.mask_paths: List[Path] = []
640
+ self.class_map: Dict[str, int] = {}
641
+
642
+ self._is_split = False
643
+ self._are_transforms_configured = False
644
+ self.train_transform: Optional[Callable] = None
645
+ self.val_transform: Optional[Callable] = None
646
+
647
+ @classmethod
648
+ def from_folders(cls, image_dir: Union[str, Path], mask_dir: Union[str, Path]) -> 'SegmentationDatasetMaker':
649
+ """
650
+ Creates a maker instance by loading all matching image-mask pairs
651
+ from two corresponding directories.
652
+
653
+ This method assumes that for an image `images/img_001.png`, there
654
+ is a corresponding mask `masks/img_001.png`.
655
+
656
+ Args:
657
+ image_dir (str | Path): Path to the directory containing input images.
658
+ mask_dir (str | Path): Path to the directory containing segmentation masks.
659
+
660
+ Returns:
661
+ SegmentationDatasetMaker: A new instance with all pairs loaded.
662
+ """
663
+ maker = cls()
664
+ img_path_obj = make_fullpath(image_dir, enforce="directory")
665
+ msk_path_obj = make_fullpath(mask_dir, enforce="directory")
666
+
667
+ # Find all images
668
+ image_files = sorted([
669
+ p for p in img_path_obj.glob('*.*')
670
+ if p.suffix.lower() in cls.IMG_EXTENSIONS
671
+ ])
672
+
673
+ if not image_files:
674
+ _LOGGER.error(f"No images with extensions {cls.IMG_EXTENSIONS} found in {image_dir}")
675
+ raise FileNotFoundError()
676
+
677
+ _LOGGER.info(f"Found {len(image_files)} images. Searching for matching masks in {mask_dir}...")
678
+
679
+ good_img_paths = []
680
+ good_mask_paths = []
681
+
682
+ for img_file in image_files:
683
+ mask_file = None
684
+
685
+ # 1. Try to find mask with the exact same name
686
+ mask_file_primary = msk_path_obj / img_file.name
687
+ if mask_file_primary.exists():
688
+ mask_file = mask_file_primary
689
+
690
+ # 2. If not, try to find mask with same stem + common mask extension
691
+ if mask_file is None:
692
+ for ext in cls.IMG_EXTENSIONS: # Masks are often .png
693
+ mask_file_secondary = msk_path_obj / (img_file.stem + ext)
694
+ if mask_file_secondary.exists():
695
+ mask_file = mask_file_secondary
696
+ break
697
+
698
+ # 3. If a match is found, add the pair
699
+ if mask_file:
700
+ good_img_paths.append(img_file)
701
+ good_mask_paths.append(mask_file)
702
+ else:
703
+ _LOGGER.warning(f"No corresponding mask found for image: {img_file.name}")
704
+
705
+ if not good_img_paths:
706
+ _LOGGER.error("No matching image-mask pairs were found.")
707
+ raise FileNotFoundError()
708
+
709
+ _LOGGER.info(f"Successfully found {len(good_img_paths)} image-mask pairs.")
710
+ maker.image_paths = good_img_paths
711
+ maker.mask_paths = good_mask_paths
712
+
713
+ return maker
714
+
715
+ @staticmethod
716
+ def inspect_folder(path: Union[str, Path]):
717
+ """
718
+ Logs a report of the types, sizes, and channels of image files
719
+ found in the directory. Useful for checking masks.
720
+ """
721
+ VisionDatasetMaker.inspect_folder(path)
722
+
723
+ def set_class_map(self, class_map: Dict[str, int]) -> 'SegmentationDatasetMaker':
724
+ """
725
+ Sets a map of pixel_value -> class_name. This is used by the MLTrainer for clear evaluation reports.
726
+
727
+ Args:
728
+ class_map (Dict[int, str]): A dictionary mapping the integer pixel
729
+ value in a mask to its string name.
730
+ Example: {'background': 0, 'road': 1, 'car': 2}
731
+ """
732
+ self.class_map = class_map
733
+ _LOGGER.info(f"Class map set: {class_map}")
734
+ return self
735
+
736
+ @property
737
+ def classes(self) -> List[str]:
738
+ """Returns the list of class names, if set."""
739
+ if self.class_map:
740
+ return list(self.class_map.keys())
741
+ return []
742
+
743
+ def split_data(self, val_size: float = 0.2, test_size: float = 0.0,
744
+ random_state: Optional[int] = 42) -> 'SegmentationDatasetMaker':
745
+ """
746
+ Splits the loaded image-mask pairs into train, validation, and test sets.
747
+
748
+ Args:
749
+ val_size (float): Proportion of the dataset to reserve for validation.
750
+ test_size (float): Proportion of the dataset to reserve for testing.
751
+ random_state (int | None): Seed for reproducible splits.
752
+
753
+ Returns:
754
+ SegmentationDatasetMaker: The same instance, now with datasets created.
755
+ """
756
+ if self._is_split:
757
+ _LOGGER.warning("Data has already been split.")
758
+ return self
759
+
760
+ if val_size + test_size >= 1.0:
761
+ _LOGGER.error("The sum of val_size and test_size must be less than 1.")
762
+ raise ValueError()
763
+
764
+ if not self.image_paths:
765
+ _LOGGER.error("There is no data to split. Use .from_folders() first.")
766
+ raise RuntimeError()
767
+
768
+ indices = list(range(len(self.image_paths)))
769
+
770
+ # Split indices
771
+ train_indices, val_test_indices = train_test_split(
772
+ indices, test_size=(val_size + test_size), random_state=random_state
773
+ )
774
+
775
+ # Helper to get paths from indices
776
+ def get_paths(idx_list):
777
+ return [self.image_paths[i] for i in idx_list], [self.mask_paths[i] for i in idx_list]
778
+
779
+ train_imgs, train_masks = get_paths(train_indices)
780
+
781
+ if test_size > 0:
782
+ val_indices, test_indices = train_test_split(
783
+ val_test_indices, test_size=(test_size / (val_size + test_size)),
784
+ random_state=random_state
785
+ )
786
+ val_imgs, val_masks = get_paths(val_indices)
787
+ test_imgs, test_masks = get_paths(test_indices)
788
+
789
+ self._test_dataset = _SegmentationDataset(test_imgs, test_masks, transform=None)
790
+ self._test_dataset.classes = self.classes # type: ignore
791
+ _LOGGER.info(f"Test set created with {len(self._test_dataset)} images.")
792
+ else:
793
+ val_imgs, val_masks = get_paths(val_test_indices)
794
+
795
+ self._train_dataset = _SegmentationDataset(train_imgs, train_masks, transform=None)
796
+ self._val_dataset = _SegmentationDataset(val_imgs, val_masks, transform=None)
797
+
798
+ # Propagate class names to datasets for MLTrainer
799
+ self._train_dataset.classes = self.classes # type: ignore
800
+ self._val_dataset.classes = self.classes # type: ignore
801
+
802
+ self._is_split = True
803
+ _LOGGER.info(f"Data split into: \n- Training: {len(self._train_dataset)} images \n- Validation: {len(self._val_dataset)} images")
804
+ return self
805
+
806
+ def configure_transforms(self,
807
+ resize_size: int = 256,
808
+ crop_size: int = 224,
809
+ mean: List[float] = [0.485, 0.456, 0.406],
810
+ std: List[float] = [0.229, 0.224, 0.225]) -> 'SegmentationDatasetMaker':
811
+ """
812
+ Configures and applies the image and mask transformations.
813
+
814
+ This method must be called AFTER data is split.
815
+
816
+ Args:
817
+ resize_size (int): The size to resize the smallest edge to
818
+ for validation/testing.
819
+ crop_size (int): The target size (square) for the final
820
+ cropped image.
821
+ mean (List[float]): The mean values for image normalization.
822
+ std (List[float]): The std dev values for image normalization.
823
+
824
+ Returns:
825
+ SegmentationDatasetMaker: The same instance, with transforms applied.
826
+ """
827
+ if not self._is_split:
828
+ _LOGGER.error("Transforms must be configured AFTER splitting data. Call .split_data() first.")
829
+ raise RuntimeError()
830
+
831
+ # --- Store components for validation recipe ---
832
+ self.val_recipe_components = {
833
+ VisionTransformRecipeKeys.RESIZE_SIZE: resize_size,
834
+ VisionTransformRecipeKeys.CROP_SIZE: crop_size,
835
+ VisionTransformRecipeKeys.MEAN: mean,
836
+ VisionTransformRecipeKeys.STD: std
837
+ }
838
+
839
+ # --- Validation/Test Pipeline (Deterministic) ---
840
+ self.val_transform = _PairedCompose([
841
+ _PairedResize(resize_size),
842
+ _PairedCenterCrop(crop_size),
843
+ _PairedToTensor(),
844
+ _PairedNormalize(mean, std)
845
+ ])
846
+
847
+ # --- Training Pipeline (Augmentation) ---
848
+ self.train_transform = _PairedCompose([
849
+ _PairedRandomResizedCrop(crop_size),
850
+ _PairedRandomHorizontalFlip(p=0.5),
851
+ _PairedToTensor(),
852
+ _PairedNormalize(mean, std)
853
+ ])
854
+
855
+ # --- Apply Transforms to the Datasets ---
856
+ self._train_dataset.transform = self.train_transform # type: ignore
857
+ self._val_dataset.transform = self.val_transform # type: ignore
858
+ if self._test_dataset:
859
+ self._test_dataset.transform = self.val_transform # type: ignore
860
+
861
+ self._are_transforms_configured = True
862
+ _LOGGER.info("Paired segmentation transforms configured and applied.")
863
+ return self
864
+
865
+ def get_datasets(self) -> Tuple[Dataset, ...]:
866
+ """
867
+ Returns the final train, validation, and optional test datasets.
868
+
869
+ Raises:
870
+ RuntimeError: If called before data is split.
871
+ RuntimeError: If called before transforms are configured.
872
+ """
873
+ if not self._is_split:
874
+ _LOGGER.error("Data has not been split. Call .split_data() first.")
875
+ raise RuntimeError()
876
+ if not self._are_transforms_configured:
877
+ _LOGGER.error("Transforms have not been configured. Call .configure_transforms() first.")
878
+ raise RuntimeError()
879
+
880
+ if self._test_dataset:
881
+ return self._train_dataset, self._val_dataset, self._test_dataset
882
+ return self._train_dataset, self._val_dataset
883
+
884
+ def save_transform_recipe(self, filepath: Union[str, Path]) -> None:
885
+ """
886
+ Saves the validation transform pipeline as a JSON recipe file.
887
+
888
+ This recipe can be loaded by the PyTorchVisionInferenceHandler
889
+ to ensure identical preprocessing.
890
+
891
+ Args:
892
+ filepath (str | Path): The path to save the .json recipe file.
893
+ """
894
+ if not self._are_transforms_configured:
895
+ _LOGGER.error("Transforms are not configured. Call .configure_transforms() first.")
896
+ raise RuntimeError()
897
+
898
+ components = self.val_recipe_components
899
+
900
+ if not components:
901
+ _LOGGER.error(f"Error getting the transformers recipe for validation set.")
902
+ raise ValueError()
903
+
904
+ # validate path
905
+ file_path = make_fullpath(filepath, make=True, enforce="file")
906
+
907
+ # Add standard transforms
908
+ recipe: Dict[str, Any] = {
909
+ VisionTransformRecipeKeys.TASK: "segmentation",
910
+ VisionTransformRecipeKeys.PIPELINE: [
911
+ {VisionTransformRecipeKeys.NAME: "Resize", "kwargs": {"size": components["resize_size"]}},
912
+ {VisionTransformRecipeKeys.NAME: "CenterCrop", "kwargs": {"size": components["crop_size"]}},
913
+ {VisionTransformRecipeKeys.NAME: "ToTensor", "kwargs": {}},
914
+ {VisionTransformRecipeKeys.NAME: "Normalize", "kwargs": {
915
+ "mean": components["mean"],
916
+ "std": components["std"]
917
+ }}
918
+ ]
919
+ }
920
+
921
+ # Save the file
922
+ save_recipe(recipe, file_path)
923
+
924
+
925
+ # Object detection
926
+ def _od_collate_fn(batch: List[Tuple[torch.Tensor, Dict[str, torch.Tensor]]]) -> Tuple[List[torch.Tensor], List[Dict[str, torch.Tensor]]]:
927
+ """
928
+ Custom collate function for object detection.
929
+
930
+ Takes a list of (image, target) tuples and zips them into two lists:
931
+ (list_of_images, list_of_targets).
932
+ This is required for models like Faster R-CNN, which accept a list
933
+ of images of varying sizes.
934
+ """
935
+ return tuple(zip(*batch)) # type: ignore
936
+
937
+
938
+ class _ObjectDetectionDataset(Dataset):
939
+ """
940
+ Internal helper class to load image-annotation pairs.
941
+
942
+ Loads an image as 'RGB' and parses its corresponding JSON annotation file
943
+ to create the required target dictionary (boxes, labels).
944
+ """
945
+ def __init__(self, image_paths: List[Path], annotation_paths: List[Path], transform: Optional[Callable] = None):
946
+ self.image_paths = image_paths
947
+ self.annotation_paths = annotation_paths
948
+ self.transform = transform
949
+
950
+ # --- Propagate 'classes' if they exist (for MLTrainer) ---
951
+ self.classes: List[str] = []
952
+
953
+ def __len__(self):
954
+ return len(self.image_paths)
955
+
956
+ def __getitem__(self, idx):
957
+ img_path = self.image_paths[idx]
958
+ ann_path = self.annotation_paths[idx]
959
+
960
+ try:
961
+ # Open image
962
+ image = Image.open(img_path).convert("RGB")
963
+
964
+ # Load and parse annotation
965
+ with open(ann_path, 'r') as f:
966
+ ann_data = json.load(f)
967
+
968
+ # Get boxes and labels from JSON
969
+ boxes = ann_data[ObjectDetectionKeys.BOXES] # Expected: [[x1, y1, x2, y2], ...]
970
+ labels = ann_data[ObjectDetectionKeys.LABELS] # Expected: [1, 2, 1, ...]
971
+
972
+ # Convert to tensors
973
+ target: Dict[str, Any] = {}
974
+ target[ObjectDetectionKeys.BOXES] = torch.as_tensor(boxes, dtype=torch.float32)
975
+ target[ObjectDetectionKeys.LABELS] = torch.as_tensor(labels, dtype=torch.int64)
976
+
977
+ except Exception as e:
978
+ _LOGGER.error(f"Error loading sample #{idx}: {img_path.name} / {ann_path.name}. Error: {e}")
979
+ # Return empty/dummy data
980
+ return torch.empty(3, 224, 224), {ObjectDetectionKeys.BOXES: torch.empty((0, 4)), ObjectDetectionKeys.LABELS: torch.empty(0, dtype=torch.long)}
981
+
982
+ if self.transform:
983
+ image, target = self.transform(image, target)
984
+
985
+ return image, target
986
+
987
+ # Internal Paired Transform Helpers for Object Detection
988
+ class _OD_PairedCompose:
989
+ """A 'Compose' for paired image/target_dict transforms."""
990
+ def __init__(self, transforms: List[Callable]):
991
+ self.transforms = transforms
992
+
993
+ def __call__(self, image: Any, target: Any) -> Tuple[Any, Any]:
994
+ for t in self.transforms:
995
+ image, target = t(image, target)
996
+ return image, target
997
+
998
+ class _OD_PairedToTensor:
999
+ """Converts a PIL Image to Tensor, passes targets dict through."""
1000
+ def __call__(self, image: Image.Image, target: Dict[str, Any]) -> Tuple[torch.Tensor, Dict[str, Any]]:
1001
+ return TF.to_tensor(image), target
1002
+
1003
+ class _OD_PairedNormalize:
1004
+ """Normalizes the image tensor and leaves the target dict untouched."""
1005
+ def __init__(self, mean: List[float], std: List[float]):
1006
+ self.normalize = transforms.Normalize(mean, std)
1007
+
1008
+ def __call__(self, image: torch.Tensor, target: Dict[str, Any]) -> Tuple[torch.Tensor, Dict[str, Any]]:
1009
+ image_normalized = self.normalize(image)
1010
+ return image_normalized, target
1011
+
1012
+ class _OD_PairedRandomHorizontalFlip:
1013
+ """Applies the same random horizontal flip to both image and targets['boxes']."""
1014
+ def __init__(self, p: float = 0.5):
1015
+ self.p = p
1016
+
1017
+ def __call__(self, image: Image.Image, target: Dict[str, Any]) -> Tuple[Image.Image, Dict[str, Any]]:
1018
+ if random.random() < self.p:
1019
+ w, h = image.size
1020
+ # Use new variable names to avoid linter confusion
1021
+ flipped_image = TF.hflip(image) # type: ignore
1022
+
1023
+ # Flip boxes
1024
+ boxes = target[ObjectDetectionKeys.BOXES].clone() # [N, 4]
1025
+
1026
+ # xmin' = w - xmax
1027
+ # xmax' = w - xmin
1028
+ boxes[:, 0] = w - target[ObjectDetectionKeys.BOXES][:, 2] # xmin'
1029
+ boxes[:, 2] = w - target[ObjectDetectionKeys.BOXES][:, 0] # xmax'
1030
+ target[ObjectDetectionKeys.BOXES] = boxes
1031
+
1032
+ return flipped_image, target # type: ignore
1033
+
1034
+ return image, target
1035
+
1036
+
1037
+ class ObjectDetectionDatasetMaker(_BaseMaker):
1038
+ """
1039
+ Creates processed PyTorch datasets for object detection from image
1040
+ and JSON annotation folders.
1041
+
1042
+ This maker finds all matching image-annotation pairs from two directories,
1043
+ splits them, and applies identical transformations (including augmentations)
1044
+ to both the image and its corresponding target dictionary.
1045
+
1046
+ The `DragonFastRCNN` model expects a list of images and a list of targets,
1047
+ so this class provides a `collate_fn` to be used with a DataLoader.
1048
+
1049
+ Workflow:
1050
+ 1. `maker = ObjectDetectionDatasetMaker.from_folders(img_dir, ann_dir)`
1051
+ 2. `maker.set_class_map({'background': 0, 'person': 1, 'car': 2})`
1052
+ 3. `maker.split_data(val_size=0.2)`
1053
+ 4. `maker.configure_transforms()`
1054
+ 5. `train_ds, val_ds = maker.get_datasets()`
1055
+ 6. `collate_fn = maker.collate_fn`
1056
+ 7. `train_loader = DataLoader(train_ds, ..., collate_fn=collate_fn)`
1057
+ """
1058
+ IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff')
1059
+
1060
+ def __init__(self):
1061
+ """
1062
+ Typically not called directly. Use the class method `from_folders()` to create an instance.
1063
+ """
1064
+ super().__init__()
1065
+ self.image_paths: List[Path] = []
1066
+ self.annotation_paths: List[Path] = []
1067
+ self.class_map: Dict[str, int] = {}
1068
+
1069
+ self._is_split = False
1070
+ self._are_transforms_configured = False
1071
+ self.train_transform: Optional[Callable] = None
1072
+ self.val_transform: Optional[Callable] = None
1073
+ self._val_recipe_components: Optional[Dict[str, Any]] = None
1074
+
1075
+ @classmethod
1076
+ def from_folders(cls, image_dir: Union[str, Path], annotation_dir: Union[str, Path]) -> 'ObjectDetectionDatasetMaker':
1077
+ """
1078
+ Creates a maker instance by loading all matching image-annotation pairs
1079
+ from two corresponding directories.
1080
+
1081
+ This method assumes that for an image `images/img_001.png`, there
1082
+ is a corresponding annotation `annotations/img_001.json`.
1083
+
1084
+ The JSON file must contain "boxes" and "labels" keys:
1085
+ `{"boxes": [[x1,y1,x2,y2], ...], "labels": [1, 2, ...]}`
1086
+
1087
+ Args:
1088
+ image_dir (str | Path): Path to the directory containing input images.
1089
+ annotation_dir (str | Path): Path to the directory containing .json
1090
+ annotation files.
1091
+
1092
+ Returns:
1093
+ ObjectDetectionDatasetMaker: A new instance with all pairs loaded.
1094
+ """
1095
+ maker = cls()
1096
+ img_path_obj = make_fullpath(image_dir, enforce="directory")
1097
+ ann_path_obj = make_fullpath(annotation_dir, enforce="directory")
1098
+
1099
+ # Find all images
1100
+ image_files = sorted([
1101
+ p for p in img_path_obj.glob('*.*')
1102
+ if p.suffix.lower() in cls.IMG_EXTENSIONS
1103
+ ])
1104
+
1105
+ if not image_files:
1106
+ _LOGGER.error(f"No images with extensions {cls.IMG_EXTENSIONS} found in {image_dir}")
1107
+ raise FileNotFoundError()
1108
+
1109
+ _LOGGER.info(f"Found {len(image_files)} images. Searching for matching .json annotations in {annotation_dir}...")
1110
+
1111
+ good_img_paths = []
1112
+ good_ann_paths = []
1113
+
1114
+ for img_file in image_files:
1115
+ # Find annotation with same stem + .json
1116
+ ann_file = ann_path_obj / (img_file.stem + ".json")
1117
+
1118
+ if ann_file.exists():
1119
+ good_img_paths.append(img_file)
1120
+ good_ann_paths.append(ann_file)
1121
+ else:
1122
+ _LOGGER.warning(f"No corresponding .json annotation found for image: {img_file.name}")
1123
+
1124
+ if not good_img_paths:
1125
+ _LOGGER.error("No matching image-annotation pairs were found.")
1126
+ raise FileNotFoundError()
1127
+
1128
+ _LOGGER.info(f"Successfully found {len(good_img_paths)} image-annotation pairs.")
1129
+ maker.image_paths = good_img_paths
1130
+ maker.annotation_paths = good_ann_paths
1131
+
1132
+ return maker
1133
+
1134
+ @staticmethod
1135
+ def inspect_folder(path: Union[str, Path]):
1136
+ """
1137
+ Logs a report of the types, sizes, and channels of image files
1138
+ found in the directory.
1139
+ """
1140
+ VisionDatasetMaker.inspect_folder(path)
1141
+
1142
+ def set_class_map(self, class_map: Dict[str, int]) -> 'ObjectDetectionDatasetMaker':
1143
+ """
1144
+ Sets a map of class_name -> pixel_value. This is used by the
1145
+ MLTrainer for clear evaluation reports.
1146
+
1147
+ **Important:** For object detection models, 'background' MUST
1148
+ be included as class 0.
1149
+ Example: `{'background': 0, 'person': 1, 'car': 2}`
1150
+
1151
+ Args:
1152
+ class_map (Dict[str, int]): A dictionary mapping the string name
1153
+ to its integer label.
1154
+ """
1155
+ if 'background' not in class_map or class_map['background'] != 0:
1156
+ _LOGGER.warning("Object detection class map should include 'background' mapped to 0.")
1157
+
1158
+ self.class_map = class_map
1159
+ _LOGGER.info(f"Class map set: {class_map}")
1160
+ return self
1161
+
1162
+ @property
1163
+ def classes(self) -> List[str]:
1164
+ """Returns the list of class names, if set."""
1165
+ if self.class_map:
1166
+ return list(self.class_map.keys())
1167
+ return []
1168
+
1169
+ def split_data(self, val_size: float = 0.2, test_size: float = 0.0,
1170
+ random_state: Optional[int] = 42) -> 'ObjectDetectionDatasetMaker':
1171
+ """
1172
+ Splits the loaded image-annotation pairs into train, validation, and test sets.
1173
+
1174
+ Args:
1175
+ val_size (float): Proportion of the dataset to reserve for validation.
1176
+ test_size (float): Proportion of the dataset to reserve for testing.
1177
+ random_state (int | None): Seed for reproducible splits.
1178
+
1179
+ Returns:
1180
+ ObjectDetectionDatasetMaker: The same instance, now with datasets created.
1181
+ """
1182
+ if self._is_split:
1183
+ _LOGGER.warning("Data has already been split.")
1184
+ return self
1185
+
1186
+ if val_size + test_size >= 1.0:
1187
+ _LOGGER.error("The sum of val_size and test_size must be less than 1.")
1188
+ raise ValueError()
1189
+
1190
+ if not self.image_paths:
1191
+ _LOGGER.error("There is no data to split. Use .from_folders() first.")
1192
+ raise RuntimeError()
1193
+
1194
+ indices = list(range(len(self.image_paths)))
1195
+
1196
+ # Split indices
1197
+ train_indices, val_test_indices = train_test_split(
1198
+ indices, test_size=(val_size + test_size), random_state=random_state
1199
+ )
1200
+
1201
+ # Helper to get paths from indices
1202
+ def get_paths(idx_list):
1203
+ return [self.image_paths[i] for i in idx_list], [self.annotation_paths[i] for i in idx_list]
1204
+
1205
+ train_imgs, train_anns = get_paths(train_indices)
1206
+
1207
+ if test_size > 0:
1208
+ val_indices, test_indices = train_test_split(
1209
+ val_test_indices, test_size=(test_size / (val_size + test_size)),
1210
+ random_state=random_state
1211
+ )
1212
+ val_imgs, val_anns = get_paths(val_indices)
1213
+ test_imgs, test_anns = get_paths(test_indices)
1214
+
1215
+ self._test_dataset = _ObjectDetectionDataset(test_imgs, test_anns, transform=None)
1216
+ self._test_dataset.classes = self.classes # type: ignore
1217
+ _LOGGER.info(f"Test set created with {len(self._test_dataset)} images.")
1218
+ else:
1219
+ val_imgs, val_anns = get_paths(val_test_indices)
1220
+
1221
+ self._train_dataset = _ObjectDetectionDataset(train_imgs, train_anns, transform=None)
1222
+ self._val_dataset = _ObjectDetectionDataset(val_imgs, val_anns, transform=None)
1223
+
1224
+ # Propagate class names to datasets for MLTrainer
1225
+ self._train_dataset.classes = self.classes # type: ignore
1226
+ self._val_dataset.classes = self.classes # type: ignore
1227
+
1228
+ self._is_split = True
1229
+ _LOGGER.info(f"Data split into: \n- Training: {len(self._train_dataset)} images \n- Validation: {len(self._val_dataset)} images")
1230
+ return self
1231
+
1232
+ def configure_transforms(self,
1233
+ mean: List[float] = [0.485, 0.456, 0.406],
1234
+ std: List[float] = [0.229, 0.224, 0.225]) -> 'ObjectDetectionDatasetMaker':
1235
+ """
1236
+ Configures and applies the image and target transformations.
1237
+
1238
+ This method must be called AFTER data is split.
1239
+
1240
+ For object detection models like Faster R-CNN, images are NOT
1241
+ resized or cropped, as the model handles variable input sizes.
1242
+ Transforms are limited to augmentation (flip), ToTensor, and Normalize.
1243
+
1244
+ Args:
1245
+ mean (List[float]): The mean values for image normalization.
1246
+ std (List[float]): The std dev values for image normalization.
1247
+
1248
+ Returns:
1249
+ ObjectDetectionDatasetMaker: The same instance, with transforms applied.
1250
+ """
1251
+ if not self._is_split:
1252
+ _LOGGER.error("Transforms must be configured AFTER splitting data. Call .split_data() first.")
1253
+ raise RuntimeError()
1254
+
1255
+ # --- Store components for validation recipe ---
1256
+ self._val_recipe_components = {
1257
+ VisionTransformRecipeKeys.MEAN: mean,
1258
+ VisionTransformRecipeKeys.STD: std
1259
+ }
1260
+
1261
+ # --- Validation/Test Pipeline (Deterministic) ---
1262
+ self.val_transform = _OD_PairedCompose([
1263
+ _OD_PairedToTensor(),
1264
+ _OD_PairedNormalize(mean, std)
1265
+ ])
1266
+
1267
+ # --- Training Pipeline (Augmentation) ---
1268
+ self.train_transform = _OD_PairedCompose([
1269
+ _OD_PairedRandomHorizontalFlip(p=0.5),
1270
+ _OD_PairedToTensor(),
1271
+ _OD_PairedNormalize(mean, std)
1272
+ ])
1273
+
1274
+ # --- Apply Transforms to the Datasets ---
1275
+ self._train_dataset.transform = self.train_transform # type: ignore
1276
+ self._val_dataset.transform = self.val_transform # type: ignore
1277
+ if self._test_dataset:
1278
+ self._test_dataset.transform = self.val_transform # type: ignore
1279
+
1280
+ self._are_transforms_configured = True
1281
+ _LOGGER.info("Paired object detection transforms configured and applied.")
1282
+ return self
1283
+
1284
+ def get_datasets(self) -> Tuple[Dataset, ...]:
1285
+ """
1286
+ Returns the final train, validation, and optional test datasets.
1287
+
1288
+ Raises:
1289
+ RuntimeError: If called before data is split.
1290
+ RuntimeError: If called before transforms are configured.
1291
+ """
1292
+ if not self._is_split:
1293
+ _LOGGER.error("Data has not been split. Call .split_data() first.")
1294
+ raise RuntimeError()
1295
+ if not self._are_transforms_configured:
1296
+ _LOGGER.error("Transforms have not been configured. Call .configure_transforms() first.")
1297
+ raise RuntimeError()
1298
+
1299
+ if self._test_dataset:
1300
+ return self._train_dataset, self._val_dataset, self._test_dataset
1301
+ return self._train_dataset, self._val_dataset
1302
+
1303
+ @property
1304
+ def collate_fn(self) -> Callable:
1305
+ """
1306
+ Returns the collate function required by a DataLoader for this
1307
+ dataset. This function ensures that images and targets are
1308
+ batched as separate lists.
1309
+ """
1310
+ return _od_collate_fn
1311
+
1312
+ def save_transform_recipe(self, filepath: Union[str, Path]) -> None:
1313
+ """
1314
+ Saves the validation transform pipeline as a JSON recipe file.
1315
+
1316
+ For object detection, this recipe only includes ToTensor and
1317
+ Normalize, as resizing is handled by the model.
1318
+
1319
+ Args:
1320
+ filepath (str | Path): The path to save the .json recipe file.
1321
+ """
1322
+ if not self._are_transforms_configured:
1323
+ _LOGGER.error("Transforms are not configured. Call .configure_transforms() first.")
1324
+ raise RuntimeError()
1325
+
1326
+ components = self._val_recipe_components
1327
+
1328
+ if not components:
1329
+ _LOGGER.error(f"Error getting the transformers recipe for validation set.")
1330
+ raise ValueError()
1331
+
1332
+ # validate path
1333
+ file_path = make_fullpath(filepath, make=True, enforce="file")
1334
+
1335
+ # Add standard transforms
1336
+ recipe: Dict[str, Any] = {
1337
+ VisionTransformRecipeKeys.TASK: "object_detection",
1338
+ VisionTransformRecipeKeys.PIPELINE: [
1339
+ {VisionTransformRecipeKeys.NAME: "ToTensor", "kwargs": {}},
1340
+ {VisionTransformRecipeKeys.NAME: "Normalize", "kwargs": {
1341
+ "mean": components["mean"],
1342
+ "std": components["std"]
1343
+ }}
1344
+ ]
1345
+ }
1346
+
1347
+ # Save the file
1348
+ save_recipe(recipe, file_path)
1349
+
1350
+
1351
+ def info():
1352
+ _script_info(__all__)