File size: 3,115 Bytes
361b108 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import os
import pandas as pd
import torch
from torch.utils.data import Dataset
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np
class NiftiSegmentationDataset(Dataset):
def __init__(self, csv_file, channel_keys = ['first_path'], augment=True):
"""
Args:
csv_file (str): CSV file with columns: 'image_path', 'label_path'
use_mask (bool): Whether to concatenate mask as extra channel (not typical for segmentation targets)
augment (bool): Apply data augmentation
"""
self.df = pd.read_csv(csv_file)
self.augment = augment
source_directory = '/workspace/data/mipsodeliadefault'
breast_mask_source_directory = '/workspace/Segmentation/breast_masks/mamamia/breastmaskmamamia'
self.samples = []
for _, row in self.df.iterrows():
image_paths = []
for key in channel_keys:
image_paths.append(os.path.join(source_directory, row['patient_id'], key + '.png'))
mask_path = os.path.join(source_directory, row['patient_id'], 'label.png')
breast_mask_path = os.path.join(breast_mask_source_directory, row['patient_id']+'.png')
self.samples.append({
"image_path": image_paths,
"breast_mask_path": breast_mask_path,
"mask_path": mask_path,
})
self.transform = self.build_transforms()
def build_transforms(self):
additional_targets = {'breast_mask': 'mask'}
if self.augment:
return A.Compose([
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.RandomRotate90(p=0.5),
A.Rotate(limit=15, p=0.5),
A.Affine(
scale=(0.95, 1.05),
shear=5,
translate_percent=(0.05, 0.05),
p=0.5
),
# A.ColorJitter(brightness=0.1, contrast=0.1, p=0.5),
A.GaussianBlur(blur_limit=3, p=0.3),
A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),
A.Normalize(mean=(0.5,), std=(0.5,)),
ToTensorV2()
], additional_targets=additional_targets)
else:
return A.Compose([
A.Normalize(mean=(0.5,), std=(0.5,)),
ToTensorV2()
], additional_targets=additional_targets)
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
sample = self.samples[idx]
# Load images as numpy arrays
image_channels = []
for image_path in sample["image_path"]:
img = np.array(Image.open(image_path).convert('L')) # Convert to grayscale
image_channels.append(img)
image = np.stack(image_channels, axis=-1) # Shape: [H, W, C]
augmented = self.transform(image=image)
image = augmented['image'] # Tensor [C, H, W]
return image |