The current dataset.py defines LensDataset and WrapperDataset to load and preprocess lens image data with optional transformations. The code mixes data loading, conditional logic for categories, and transformations in a single class, making it less flexible and harder to extend. The transform pipeline in get_transforms is partially commented out and lacks clear augmentation handling. The issue is to refactor these classes for modularity, support dynamic augmentations, remove hardcoded category logic, and improve readability and maintainability of the dataset pipeline.
The current dataset.py defines LensDataset and WrapperDataset to load and preprocess lens image data with optional transformations. The code mixes data loading, conditional logic for categories, and transformations in a single class, making it less flexible and harder to extend. The transform pipeline in get_transforms is partially commented out and lacks clear augmentation handling. The issue is to refactor these classes for modularity, support dynamic augmentations, remove hardcoded category logic, and improve readability and maintainability of the dataset pipeline.