import os import cv2 from matplotlib import pyplot as plt import pandas as pd import numpy as np import torch from torch.utils.data import Dataset import random import os import cv2 import pandas as pd import torch from torch.utils.data import Dataset import os import cv2 import torch import numpy as np import pandas as pd from torch.utils.data import Dataset import SimpleITK as sitk import matplotlib.pyplot as plt import os class BreastMRI_ABMIL_NII(Dataset): def __init__(self, csv_file = "/anvme/workspace/b268dc11-breastt/BreastMRI/DataExtraction/combined_odelia_with_patient_split.csv", root_dir = "/anvme/workspace/b268dc11-breastt/OdeliaV2_extracted/OdeliaV2", split='train', transform=None): self.data = pd.read_csv(csv_file) print("Available splits:", self.data['Split_patient'].value_counts()) if split is not None: self.data = self.data[self.data['Split_patient'] == split].reset_index(drop=True) print(f"Split = {split} -> {len(self.data)} samples") self.transform = transform self.root_dir = root_dir def __len__(self): return len(self.data) def process(self, dyn0, dyn1): sub = dyn1-dyn0 sub = sub-sub.min() # Note: negative values causes overflow when using uint sub = sub.astype(np.uint16) return sub def save_middle_slices_as_png(self, tensor, output_dir, prefix="sample"): os.makedirs(output_dir, exist_ok=True) channel_names = ['Pre', 'Post', 'Subtraction'] # If batch dimension exists, take first sample if tensor.ndim == 4: tensor = tensor[0] # shape (3, 224, 224) for i in range(3): slice_img = tensor[i].cpu().numpy() # entire 2D image per channel plt.figure(figsize=(6,6)) plt.imshow(slice_img, cmap='gray') plt.axis('off') filename = os.path.join(output_dir, f"{prefix}_{channel_names[i]}.png") plt.savefig(filename, bbox_inches='tight', pad_inches=0) plt.close() def resize(self, arr): # Initialize output array with target shape and same dtype arr_resized = np.zeros((arr.shape[0], 224, 224), dtype=arr.dtype) # Resize each slice individually for i in range(arr.shape[0]): arr_resized[i] = cv2.resize(arr[i], (224, 224), interpolation=cv2.INTER_LINEAR) return arr_resized def normalize_per_channel_per_sample(self, x): # Compute min and max per sample and channel: shape (batch_size, channels, 1, 1) mins = x.amin(dim=[2,3], keepdim=True) maxs = x.amax(dim=[2,3], keepdim=True) # Avoid division by zero (if max == min) denom = (maxs - mins).clamp(min=1e-6) # Normalize x_norm = (x - mins) / denom return x_norm def _load_slices(self, row_value): filename = os.path.join(self.root_dir, row_value.replace('\\', '/')) # Try posts in descending order post_files = ["Post_4.nii.gz", "Post_3.nii.gz", "Post_2.nii.gz"] post_path = None for pf in post_files: candidate = os.path.join(filename, pf) if os.path.exists(candidate): post_path = candidate break if post_path is None: raise FileNotFoundError(f"No Post_* file found in {filename}") pre_path = os.path.join(filename, "Pre.nii.gz") post_1_path = os.path.join(filename, "Post_1.nii.gz") pre_sequence = sitk.ReadImage(pre_path) post_sequence = sitk.ReadImage(post_path) post_1_sequence = sitk.ReadImage(post_1_path) pre_array = sitk.GetArrayFromImage(pre_sequence) post_1_array = sitk.GetArrayFromImage(post_1_sequence) post_array = sitk.GetArrayFromImage(post_sequence) sub_array = self.process(pre_array, post_1_array) pre_array = self.resize(pre_array) post_array = self.resize(post_array) sub_array = self.resize(sub_array) stacked_array = np.stack([pre_array, post_array, sub_array], axis=1) stacked_tensor = torch.from_numpy(stacked_array).float() normalized_tensor = self.normalize_per_channel_per_sample(stacked_tensor) return normalized_tensor def __getitem__(self, idx): if idx >= len(self.data): raise IndexError(f"[{self.__class__.__name__}] Index {idx} is out of bounds for dataset of length {len(self.data)}") row = self.data.iloc[idx] path = row['SamplePath'] stacked_tensor = self._load_slices(path) #if idx<3: # self.save_middle_slices_as_png(stacked_tensor, output_dir="./visualizations", prefix=f"sample_{idx}") label = torch.tensor(row['Lesion']).long() return stacked_tensor, label class BreastMRI_ABMIL(Dataset): """ PyTorch Dataset for patient-level ABMIL training using 32 slices per patient. Expects a CSV with columns: - 'UID' (unique per patient) - 'SamplePath' (relative path from root_dir) - 'Lesion' (patient-level label) Expects two root folders: - One for pre/post slices: root_dir_prepost - One for sub slices: root_dir_sub For each UID, constructs 32 slices: - Each slice is a 3-channel image: [pre, post, sub] - Returns a tuple: (tensor of shape [32, 3, 224, 224], label) """ def __init__(self, csv_file, root_dir_prepost, root_dir_sub, split='train', transform=None): self.data = pd.read_csv(csv_file) print("Available splits:", self.data['Split'].value_counts()) if split is not None: self.data = self.data[self.data['Split'] == split].reset_index(drop=True) print(f"Split = {split} -> {len(self.data)} samples") self.root_dir_prepost = root_dir_prepost self.root_dir_sub = root_dir_sub self.transform = transform def __len__(self): return len(self.data) def _load_slice_triplet(self, pre_path, post_path, sub_path): def read_img(p): img = cv2.imread(p, cv2.IMREAD_GRAYSCALE) img = cv2.resize(img, (224, 224)) return torch.from_numpy(img).float() / 255.0 # (H, W) pre = read_img(pre_path) post = read_img(post_path) sub = read_img(sub_path) stacked = torch.stack([pre, post, sub], dim=0) # (3, H, W) if self.transform: stacked = self.transform(stacked) return stacked # (3, H, W) def __getitem__(self, idx): if idx >= len(self.data): raise IndexError(f"[{self.__class__.__name__}] Index {idx} is out of bounds for dataset of length {len(self.data)}") row = self.data.iloc[idx] uid = row['UID'] sample_path = row['SamplePath'].replace("\\", "/") # normalize prepost_dir = os.path.join(self.root_dir_prepost, sample_path) sub_dir = os.path.join(self.root_dir_sub, sample_path) slice_stack = [] for i in range(32): slice_id = f"{i:03d}" pre_path = os.path.join(prepost_dir, f"{uid}_pre_slice_{slice_id}.png") post_path = os.path.join(prepost_dir, f"{uid}_post_slice_{slice_id}.png") sub_path = os.path.join(sub_dir, f"{uid}_sub_slice_{slice_id}.png") try: slice_img = self._load_slice_triplet(pre_path, post_path, sub_path) except Exception as e: raise FileNotFoundError(f"Error loading slice {i} for UID {uid}: {e}") slice_stack.append(slice_img) volume = torch.stack(slice_stack, dim=0) # (32, 3, 224, 224) label = torch.tensor(row['Lesion']).long() return volume, label class BreastMRI(Dataset): """ PyTorch Dataset for breast MRI slices with pre-contrast, post-contrast, and subtraction images. Expects a CSV with columns 'Split', 'pre_png', 'post_png', 'sub_png', and 'Lesion'. Paths in the CSV may be absolute or relative to `root_dir`. Returns a dict with keys 'pre', 'post', 'sub', and 'label'. """ def __init__(self, csv_file, root_dir=None, split='train', transform=None): # Load the annotations CSV self.data = pd.read_csv(csv_file) print("Available splits:", self.data['Split'].value_counts()) if split is not None: self.data = self.data[self.data['Split'] == split].reset_index(drop=True) print(f"Split = {split} -> {len(self.data)} samples") # Filter by split if provided self.root_dir = root_dir self.transform = transform def __len__(self): return len(self.data) def _load_image(self, img_path): # Read as grayscale (2D) image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) image = cv2.resize(image, (224,224)) image = np.stack([image] * 3, axis=-1) if image is None: raise FileNotFoundError(f"Could not load image at {img_path}") # Optionally apply transform if self.transform: # Assume transform can handle single-channel images image = self.transform(image) else: # Convert to tensor with shape (1, H, W) image = torch.from_numpy(image).unsqueeze(0).float() return image def __getitem__(self, idx): if idx >= len(self.data): raise IndexError(f"[{self.__class__.__name__}] Index {idx} is out of bounds for dataset of length {len(self.data)}") row = self.data.iloc[idx] try: pre_img = self._load_image(row['pre_png']) post_img = self._load_image(row['post_png']) sub_img = self._load_image(row['sub_png']) except Exception as e: raise FileNotFoundError(f"Could not load one of the images: {e}\nRow:\n{row}") label = torch.tensor(row['Lesion']).long() return { 'pre': pre_img, 'post': post_img, 'sub': sub_img, 'label': label } class MammographyDataset(Dataset): def __init__(self, split_file, data_folder, task, view, probability, split, attention_head, transform=None, rage = True): self.mammography = pd.read_csv(split_file) if view in ['CC', 'MLO']: self.mammography = self.mammography[self.mammography['View Position'] == view] self.data_folder = data_folder self.task = task self.transform = transform self.probability = probability self.split = split self.attention_head = attention_head self.rage = rage def __len__(self): return len(self.mammography) def class_name_to_labels(self, idx): if self.task == 'BreastDensity': class_name = self.mammography.iloc[idx, 8] if class_name in ['DENSITY A']: labels = 0.0 elif class_name in ['DENSITY B']: labels = 1.0 elif class_name in ['DENSITY C']: labels = 2.0 elif class_name in ['DENSITY D']: labels = 3.0 return labels elif self.task == 'BIRADS': class_name = self.mammography.iloc[idx, 7] if class_name in ['BI-RADS 1']: labels = 0.0 elif class_name in ['BI-RADS 2']: labels = 1.0 elif class_name in ['BI-RADS 3']: labels = 2.0 elif class_name in ['BI-RADS 4']: labels = 3.0 elif class_name in ['BI-RADS 5']: labels = 4.0 return labels elif self.task == 'MassRest': class_name = self.mammography.iloc[idx, 30] if 'Mass' in class_name: labels = 0.0 else: labels = 1.0 return labels elif self.task == 'MassNormal': class_name = self.mammography.iloc[idx, 30] if 'Mass' in class_name: labels = 0.0 elif 'No Finding' in class_name: labels = 1.0 else: labels = 2.0 return labels elif self.task == 'CMMD_Malignant': class_name = self.mammography.iloc[idx, 4] if class_name in ['Benign']: labels = 0.0 elif class_name in ['Malignant']: labels = 1.0 return labels elif self.task == 'CM': class_name = self.mammography.iloc[idx, 11] if class_name in ['Normal']: labels = 0.0 elif class_name in ['Benign']: labels = 1.0 elif class_name in ['Malignant']: labels = 2.0 return labels elif self.task == 'CM_Malignant': class_name = self.mammography.iloc[idx, 11] if class_name in ['Normal']: labels = 0.0 elif class_name in ['Benign']: labels = 1.0 elif class_name in ['Malignant']: labels = 2.0 return labels elif self.task in ['Calc', 'Mass']: class_name = self.mammography.iloc[idx, 9] if class_name in ['BENIGN_WITHOUT_CALLBACK']: labels = 0.0 elif class_name in ['BENIGN']: labels = 1.0 elif class_name in ['MALIGNANT']: labels = 2.0 return labels def image_load(self, idx, img_column, folder_column): if self.task in ['BreastDensity', 'MassRest', 'MassNormal']: img_name = self.mammography.iloc[idx, img_column] img_folder = self.mammography.iloc[idx, folder_column] image_name = os.path.join(self.data_folder, img_folder, img_name, "img.png") elif self.task in ['CM', 'Calc', 'Mass', 'CMMD_Malignant', 'CM_Malignant']: img_filename = self.mammography.iloc[idx, img_column] img_folder = img_filename.split("/")[-1][:-4] image_name = os.path.join(self.data_folder, img_folder, "img.png") #print("Image Name", image_name) #image_name = image_name + '/img.png' image = cv2.imread(image_name) image[image < 40] = 0 # had it for breast density, why ? image = cv2.resize(image, (448,448)) return image def convert_to_mask(self, image, threshold): gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) # Apply binary thresholding _, binary_mask = cv2.threshold(gray_image, threshold, 255, cv2.THRESH_BINARY) return binary_mask def mask_load(self, idx, img_column, folder_column): if self.task in ['BreastDensity', 'MassRest', 'MassNormal']: img_name = self.mammography.iloc[idx, img_column] img_folder = self.mammography.iloc[idx, folder_column] if self.rage: self.attention_head=random.randint(0,5) if self.attention_head in ['0', 0]: image_name = os.path.join(self.data_folder, img_folder, img_name, "attn-head0.png") if self.attention_head in ['1', 1]: image_name = os.path.join(self.data_folder, img_folder, img_name, "attn-head1.png") if self.attention_head in ['2', 2]: image_name = os.path.join(self.data_folder, img_folder, img_name, "attn-head2.png") if self.attention_head in ['3', 3]: image_name = os.path.join(self.data_folder, img_folder, img_name, "attn-head3.png") if self.attention_head in ['4', 4]: image_name = os.path.join(self.data_folder, img_folder, img_name, "attn-head4.png") if self.attention_head in ['5', 5]: image_name = os.path.join(self.data_folder, img_folder, img_name, "attn-head5.png") elif self.task in ['CM', 'Calc', 'Mass', 'CMMD_Malignant', 'CM_Malignant']: img_filename = self.mammography.iloc[idx, img_column] img_folder = img_filename.split("/")[-1][:-4] if self.rage: self.attention_head=random.randint(0,5) if self.attention_head in ['0', 0]: image_name = os.path.join(self.data_folder, img_folder, "attn-head0.png") if self.attention_head in ['1', 1]: image_name = os.path.join(self.data_folder, img_folder, "attn-head1.png") if self.attention_head in ['2', 2]: image_name = os.path.join(self.data_folder, img_folder, "attn-head2.png") if self.attention_head in ['3', 3]: image_name = os.path.join(self.data_folder, img_folder, "attn-head3.png") if self.attention_head in ['4', 4]: image_name = os.path.join(self.data_folder, img_folder, "attn-head4.png") if self.attention_head in ['5', 5]: image_name = os.path.join(self.data_folder, img_folder, "attn-head5.png") image = cv2.imread(image_name) image = self.convert_to_mask(image, threshold = 45) image = cv2.resize(image, (448,448)) return image def __getitem__(self, idx): if torch.is_tensor(idx): idx = idx.tolist() if self.task in ['BreastDensity', 'MassRest', 'MassNormal']: image = self.image_load(idx, 2, 0) mask = self.mask_load(idx, 2, 0) elif self.task in ['CM']: image = self.image_load(idx, 15, 18) mask = self.mask_load(idx, 15, 18) elif self.task in ['CM_Malignant']: image = self.image_load(idx, 12, 18) mask = self.mask_load(idx, 12, 18) elif self.task in ['Calc','Mass']: image = self.image_load(idx, 13, 0) mask = self.mask_load(idx, 13, 0) elif self.task in ['CMMD_Malignant']: image = self.image_load(idx, 9, 0) mask = self.mask_load(idx, 9, 0) if (torch.rand(1) < self.probability) and self.split == 'train' and self.task in ['CM', 'BreastDensity', 'MassRest', 'MassNormal']: mask = cv2.bitwise_not(mask) image = cv2.bitwise_and(image, image, mask=mask) elif (torch.rand(1) < self.probability) and self.split == 'train' and self.task in ['Calc', 'Mass', 'CMMD_Malignant', 'CM_Malignant']: #mask = cv2.bitwise_not(mask) image = cv2.bitwise_and(image, image, mask=mask) image = torch.from_numpy(image) image = image.permute(2, 0, 1) if self.transform: image = self.transform(image) label = self.class_name_to_labels(idx) labels = torch.from_numpy(np.array(label)) return image, labels class MammographyDatasetViz(Dataset): def __init__(self, split_file, data_folder, task, transform=None): self.mammography = pd.read_csv(split_file) self.data_folder = data_folder self.task = task self.transform = transform def __len__(self): return len(self.mammography) def class_name_to_labels(self, idx): if self.task == 'BreastDensity': class_name = self.mammography.iloc[idx, 8] if class_name in ['DENSITY A']: labels = 0.0 elif class_name in ['DENSITY B']: labels = 1.0 elif class_name in ['DENSITY C']: labels = 2.0 elif class_name in ['DENSITY D']: labels = 3.0 return labels elif self.task == 'BIRADS': class_name = self.mammography.iloc[idx, 7] if class_name in ['BI-RADS 1']: labels = 0.0 elif class_name in ['BI-RADS 2']: labels = 1.0 elif class_name in ['BI-RADS 3']: labels = 2.0 elif class_name in ['BI-RADS 4']: labels = 3.0 elif class_name in ['BI-RADS 5']: labels = 4.0 return labels def image_load(self, idx, img_column, folder_column): img_name = self.mammography.iloc[idx, img_column] study_id = self.mammography.iloc[idx, folder_column] image_name = os.path.join(self.data_folder, study_id, img_name) image_name = image_name + '.png' image = cv2.imread(image_name) #image = cv2.imresize(image, (512,512)) image = torch.from_numpy(image) #image = image.transpose(2, 0 , 1) image = image.permute(2, 0, 1) return image, study_id, img_name def __getitem__(self, idx): if torch.is_tensor(idx): idx = idx.tolist() image, study_id, img_name = self.image_load(idx, 2, 0) # 0 image folder, 1, 2nd column:image name #plt.imshow(image[0]) #plt.savefig('image.png') if self.transform: image = self.transform(image) #plt.imshow(image[0][0]) #plt.savefig('transformimage.png') label = self.class_name_to_labels(idx) labels = torch.from_numpy(np.array(label)) return image, labels, study_id, img_name ''' Example usage train_file = 'train.csv' valid_file = 'validation.csv' test_file = 'test.csv' data_folder = '/cluster/eq27ifuw/dataset_png' task = 'BreastDensity' ## Alternatively 'BIRADS' train_transforms = torch.transforms(........) test_transforms = torch.transforms(........) train_dataset = MammographyDataset(split_file = train_file, data_folder = data_folder, task = task, transform=train_transforms) valid_dataset = MammographyDataset(split_file = valid_file, data_folder = data_folder, task = task, transform=test_transforms) test_dataset = MammographyDataset(split_file = test_file, data_folder = data_folder, task = task, transform=test_transforms) '''