Spaces:
Runtime error
Runtime error
| """ | |
| Dataset Loader for BraTS 2021 PNG Dataset | |
| Handles loading of FLAIR, RGB images and segmentation masks | |
| """ | |
| import os | |
| import numpy as np | |
| import pandas as pd | |
| from PIL import Image | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader | |
| from typing import Tuple, Dict, List, Optional | |
| import albumentations as A | |
| from albumentations.pytorch import ToTensorV2 | |
| class BraTSDataset(Dataset): | |
| """BraTS 2021 Dataset Loader for PNG images""" | |
| def __init__( | |
| self, | |
| data_path: str, | |
| manifest_csv: Optional[str] = None, | |
| image_size: int = 256, | |
| mode: str = 'train', | |
| transform: Optional[A.Compose] = None, | |
| use_cache: bool = False | |
| ): | |
| """ | |
| Args: | |
| data_path: Path to brats_png directory | |
| manifest_csv: Path to CSV manifest file | |
| image_size: Target image size | |
| mode: 'train', 'val', or 'test' | |
| transform: Albumentations transform | |
| use_cache: Cache images in memory | |
| """ | |
| self.data_path = data_path | |
| self.image_size = image_size | |
| self.mode = mode | |
| self.transform = transform | |
| self.use_cache = use_cache | |
| self.cache = {} | |
| # Load data paths | |
| self.data_list = self._load_data_list(manifest_csv) | |
| def _load_data_list(self, manifest_csv: Optional[str]) -> List[Dict]: | |
| """Load list of image paths""" | |
| data_list = [] | |
| if manifest_csv and os.path.exists(manifest_csv): | |
| # Load from CSV manifest | |
| df = pd.read_csv(manifest_csv) | |
| data_list = df.to_dict('records') | |
| else: | |
| # Scan directory structure | |
| for patient_folder in sorted(os.listdir(self.data_path)): | |
| patient_path = os.path.join(self.data_path, patient_folder) | |
| if not os.path.isdir(patient_path): | |
| continue | |
| # Get all slices for this patient | |
| files = os.listdir(patient_path) | |
| # Group by slice number | |
| slices = {} | |
| for f in files: | |
| if '_FLAIR.png' in f: | |
| slice_num = f.split('_slice_')[1].split('_')[0] | |
| slices[slice_num] = { | |
| 'patient_id': patient_folder, | |
| 'slice_num': slice_num | |
| } | |
| # Add full paths | |
| for slice_num, info in slices.items(): | |
| base_name = f"{patient_folder}_slice_{slice_num}" | |
| data_list.append({ | |
| 'patient_id': patient_folder, | |
| 'slice_num': slice_num, | |
| 'flair_path': os.path.join(patient_path, f"{base_name}_FLAIR.png"), | |
| 'rgb_path': os.path.join(patient_path, f"{base_name}_RGB.png"), | |
| 'et_mask_path': os.path.join(patient_path, f"{base_name}_ET.png"), | |
| 'tc_mask_path': os.path.join(patient_path, f"{base_name}_TC.png"), | |
| 'wt_mask_path': os.path.join(patient_path, f"{base_name}_WT.png"), | |
| }) | |
| return data_list | |
| def __len__(self) -> int: | |
| return len(self.data_list) | |
| def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: | |
| """Get a single sample""" | |
| # Check cache | |
| if self.use_cache and idx in self.cache: | |
| return self.cache[idx] | |
| item = self.data_list[idx] | |
| # Load images | |
| flair = np.array(Image.open(item['flair_path']).convert('L')) | |
| rgb = np.array(Image.open(item['rgb_path']).convert('RGB')) | |
| # Load masks | |
| et_mask = np.array(Image.open(item['et_mask_path']).convert('L')) | |
| tc_mask = np.array(Image.open(item['tc_mask_path']).convert('L')) | |
| wt_mask = np.array(Image.open(item['wt_mask_path']).convert('L')) | |
| # Combine masks into multi-channel | |
| mask = np.stack([et_mask, tc_mask, wt_mask], axis=-1) | |
| # Normalize masks to 0-1 | |
| mask = (mask > 0).astype(np.float32) | |
| # Apply transforms | |
| if self.transform: | |
| transformed = self.transform(image=rgb, mask=mask) | |
| rgb = transformed['image'] | |
| mask = transformed['mask'] | |
| else: | |
| # Default: resize and convert to tensor | |
| rgb = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0 | |
| mask = torch.from_numpy(mask).permute(2, 0, 1).float() | |
| sample = { | |
| 'image': rgb, | |
| 'mask': mask, | |
| 'patient_id': item['patient_id'], | |
| 'slice_num': item['slice_num'], | |
| 'flair': torch.from_numpy(flair).unsqueeze(0).float() / 255.0 | |
| } | |
| # Cache if enabled | |
| if self.use_cache: | |
| self.cache[idx] = sample | |
| return sample | |
| def get_transforms(image_size: int = 256, mode: str = 'train') -> A.Compose: | |
| """Get augmentation transforms""" | |
| if mode == 'train': | |
| return A.Compose([ | |
| A.Resize(image_size, image_size), | |
| A.HorizontalFlip(p=0.5), | |
| A.VerticalFlip(p=0.5), | |
| A.RandomRotate90(p=0.5), | |
| A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, p=0.5), | |
| A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5), | |
| A.GaussNoise(var_limit=(10.0, 50.0), p=0.3), | |
| A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ToTensorV2() | |
| ]) | |
| else: | |
| return A.Compose([ | |
| A.Resize(image_size, image_size), | |
| A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ToTensorV2() | |
| ]) | |
| def create_dataloaders( | |
| data_path: str, | |
| batch_size: int = 8, | |
| image_size: int = 256, | |
| num_workers: int = 4, | |
| train_split: float = 0.7, | |
| val_split: float = 0.15, | |
| test_split: float = 0.15 | |
| ) -> Tuple[DataLoader, DataLoader, DataLoader]: | |
| """Create train, validation, and test dataloaders""" | |
| # Create full dataset | |
| full_dataset = BraTSDataset( | |
| data_path=data_path, | |
| image_size=image_size, | |
| mode='train' | |
| ) | |
| # Split dataset | |
| total_size = len(full_dataset) | |
| train_size = int(total_size * train_split) | |
| val_size = int(total_size * val_split) | |
| test_size = total_size - train_size - val_size | |
| train_dataset, val_dataset, test_dataset = torch.utils.data.random_split( | |
| full_dataset, | |
| [train_size, val_size, test_size], | |
| generator=torch.Generator().manual_seed(42) | |
| ) | |
| # Apply transforms | |
| train_dataset.dataset.transform = get_transforms(image_size, 'train') | |
| val_dataset.dataset.transform = get_transforms(image_size, 'val') | |
| test_dataset.dataset.transform = get_transforms(image_size, 'test') | |
| # Create dataloaders | |
| train_loader = DataLoader( | |
| train_dataset, | |
| batch_size=batch_size, | |
| shuffle=True, | |
| num_workers=num_workers, | |
| pin_memory=True | |
| ) | |
| val_loader = DataLoader( | |
| val_dataset, | |
| batch_size=batch_size, | |
| shuffle=False, | |
| num_workers=num_workers, | |
| pin_memory=True | |
| ) | |
| test_loader = DataLoader( | |
| test_dataset, | |
| batch_size=batch_size, | |
| shuffle=False, | |
| num_workers=num_workers, | |
| pin_memory=True | |
| ) | |
| return train_loader, val_loader, test_loader | |
| if __name__ == "__main__": | |
| # Test dataset loader | |
| data_path = "D:/major projrct PNG folder/brats_png" | |
| dataset = BraTSDataset(data_path=data_path, image_size=256) | |
| print(f"Dataset size: {len(dataset)}") | |
| sample = dataset[0] | |
| print(f"Image shape: {sample['image'].shape}") | |
| print(f"Mask shape: {sample['mask'].shape}") | |
| print(f"Patient ID: {sample['patient_id']}") | |
| print(f"Slice number: {sample['slice_num']}") | |