| import os | |
| import cv2 | |
| import numpy as np | |
| from torch.utils.data import Dataset as BaseDataset | |
| class SegmentationDataset(BaseDataset): | |
| """Dataset class for semantic segmentation task.""" | |
| def __init__(self, data_dir, classes=['background', 'object'], | |
| augmentation=None, preprocessing=None): | |
| self.image_dir = os.path.join(data_dir, 'Images') | |
| self.mask_dir = os.path.join(data_dir, 'Masks') | |
| for dir_path in [self.image_dir, self.mask_dir]: | |
| if not os.path.exists(dir_path): | |
| raise FileNotFoundError(f"Directory not found: {dir_path}") | |
| self.filenames = self._get_filenames() | |
| self.image_paths = [os.path.join(self.image_dir, fname) for fname in self.filenames] | |
| self.mask_paths = self._get_mask_paths() | |
| self.target_classes = [cls for cls in classes if cls.lower() != 'background'] | |
| self.class_values = [i for i, cls in enumerate(classes) if cls.lower() != 'background'] | |
| self.augmentation = augmentation | |
| self.preprocessing = preprocessing | |
| def __getitem__(self, index): | |
| image = self._load_image(self.image_paths[index]) | |
| mask = self._load_mask(self.mask_paths[index]) | |
| if self.augmentation: | |
| processed = self.augmentation(image=image, mask=mask) | |
| image, mask = processed['image'], processed['mask'] | |
| if self.preprocessing: | |
| processed = self.preprocessing(image=image, mask=mask) | |
| image, mask = processed['image'], processed['mask'] | |
| return image, mask | |
| def __len__(self): | |
| return len(self.filenames) | |
| def _get_filenames(self): | |
| """Returns sorted list of filenames, excluding directories.""" | |
| files = sorted(os.listdir(self.image_dir)) | |
| return [f for f in files if not os.path.isdir(os.path.join(self.image_dir, f))] | |
| def _get_mask_paths(self): | |
| """Generates corresponding mask paths for each image.""" | |
| mask_paths = [] | |
| for image_file in self.filenames: | |
| name, _ = os.path.splitext(image_file) | |
| mask_paths.append(os.path.join(self.mask_dir, f"{name}.png")) | |
| return mask_paths | |
| def _load_image(self, path): | |
| """Loads and converts image to RGB.""" | |
| image = cv2.imread(path) | |
| return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| def _load_mask(self, path): | |
| """Loads and processes segmentation mask.""" | |
| mask = cv2.imread(path, 0) | |
| masks = [(mask == value) for value in self.class_values] | |
| mask = np.stack(masks, axis=-1).astype('float') | |
| return mask | |
| class InferenceDataset(BaseDataset): | |
| """Dataset class for inference without ground truth masks.""" | |
| def __init__(self, data_dir, classes=['background', 'object'], | |
| augmentation=None, preprocessing=None): | |
| self.filenames = sorted([ | |
| f for f in os.listdir(data_dir) | |
| if not os.path.isdir(os.path.join(data_dir, f)) | |
| ]) | |
| self.image_paths = [os.path.join(data_dir, fname) for fname in self.filenames] | |
| self.target_classes = [cls for cls in classes if cls.lower() != 'background'] | |
| self.class_values = [i for i, cls in enumerate(classes) if cls.lower() != 'background'] | |
| self.augmentation = augmentation | |
| self.preprocessing = preprocessing | |
| def __getitem__(self, index): | |
| image = cv2.imread(self.image_paths[index]) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| original_height, original_width = image.shape[:2] | |
| if self.augmentation: | |
| image = self.augmentation(image=image)['image'] | |
| if self.preprocessing: | |
| image = self.preprocessing(image=image)['image'] | |
| return image, original_height, original_width | |
| def __len__(self): | |
| return len(self.filenames) | |
| class StreamingDataset(BaseDataset): | |
| """Dataset class optimized for video frame processing.""" | |
| def __init__(self, data_dir, classes=['background', 'object'], | |
| augmentation=None, preprocessing=None): | |
| self.filenames = self._get_frame_filenames(data_dir) | |
| self.image_paths = [os.path.join(data_dir, fname) for fname in self.filenames] | |
| self.target_classes = [cls for cls in classes if cls.lower() != 'background'] | |
| self.class_values = [i for i, cls in enumerate(classes) if cls.lower() != 'background'] | |
| self.augmentation = augmentation | |
| self.preprocessing = preprocessing | |
| def _get_frame_filenames(self, directory): | |
| """Returns sorted list of frame filenames.""" | |
| files = sorted(os.listdir(directory)) | |
| return [f for f in files if (('frame' in f or 'Image' in f) and | |
| f.lower().endswith('jpg') and | |
| not os.path.isdir(os.path.join(directory, f)))] | |
| def __getitem__(self, index): | |
| return InferenceDataset.__getitem__(self, index) | |
| def __len__(self): | |
| return len(self.filenames) |