import torchio as tio import torch import numpy as np import os import re class PreprocessDataset: def __init__(self, two_dimension = False): self.two_dimension = two_dimension def process(self, path_dir): # Implement the preprocessing logic here post_folders = [] # Extract timepoint index from folder names for foldername in os.listdir(path_dir): # match = re.match(r"dce-breast-mri-timepoint-(\d+)", foldername) #Careful: the timepoint folders in the test input are not named like this match = re.match(r"dce-breast-mri-(\d+)", foldername) if match: index = int(match.group(1)) post_folders.append((index, foldername)) # Sort folders by timepoint index post_folders_sorted = sorted(post_folders, key=lambda x: x[0]) # Extract first, second, and last (if enough exist) pre = post_folders_sorted[0][1] if len(post_folders_sorted) >= 1 else None first = post_folders_sorted[1][1] if len(post_folders_sorted) >= 2 else None last = post_folders_sorted[-1][1] if len(post_folders_sorted) >= 1 else None print(f"First: {pre}") print(f"Second: {first}") print(f"Last: {last}") full_paths = [] for p in [pre, first, last]: if p is not None: folder_path = os.path.join(path_dir, p) mha_files = [f for f in os.listdir(folder_path) if f.endswith('.mha')] if mha_files: full_paths.append(os.path.join(folder_path, mha_files[0])) # You can select another strategy if multiple .mha files exist else: print(f"No .mha files found in {folder_path}") left_array, right_array = preprocess_single(full_paths) if self.two_dimension: left_array = normalize(compute_mip(left_array)) right_array = normalize(compute_mip(right_array)) return left_array, right_array def process_self(self, path_dir): # Implement the preprocessing logic here post_folders = [] # Extract timepoint index from folder names for file in os.listdir(path_dir): # match = re.match(r"dce-breast-mri-timepoint-(\d+)", foldername) #Careful: the timepoint folders in the test input are not named like this match = re.match(r"Post_(\d+)\.nii\.gz", file) if match: index = int(match.group(1)) post_folders.append((index, file)) # Sort folders by timepoint index post_folders_sorted = sorted(post_folders, key=lambda x: x[0]) # Extract first, second, and last (if enough exist) # t2 = os.path.join(path_dir, 'T2.nii.gz') pre = os.path.join(path_dir, 'Pre.nii.gz') first = post_folders_sorted[0][1] if len(post_folders_sorted) >= 1 else None second = post_folders_sorted[1][1] if len(post_folders_sorted) >= 2 else None last = post_folders_sorted[-1][1] if len(post_folders_sorted) >= 1 else None print(f"Pre: {pre}") print(f"First: {first}") print(f"Second: {second}") print(f"Last: {last}") full_paths = [] for p in [pre, first, second, last]: if p is not None: folder_path = os.path.join(path_dir, p) full_paths.append(folder_path) # You can select another strategy if multiple .mha files exist left_array, right_array = preprocess_single(full_paths) if self.two_dimension: left_array = np.stack([normalize(compute_mip(left_array[..., i])) for i in range(left_array.shape[-1])], axis=-1) right_array = np.stack([normalize(compute_mip(right_array[..., i])) for i in range(right_array.shape[-1])], axis=-1) def compute_mip(volume): return np.max(volume, axis=2) def compute_subtraction(post, pre): return np.clip(post - pre, a_min=0, a_max=None) def normalize(mip): mip = (mip - np.min(mip)) / (np.ptp(mip)) * 255 mip = mip.astype(np.uint8) def crop_breast_height(image, margin_top=10): "Crop height to 256 and try to cover breast based on intensity localization" threshold = int(np.quantile(image.data.float(), 0.9)) foreground = image.data>threshold fg_rows = foreground[0].sum(axis=(0, 2)) top = min(max(512-int(torch.argwhere(fg_rows).max()) - margin_top, 0), 256) bottom = 256-top return tio.Crop((0,0, bottom, top, 0, 0)) def preprocess_single(paths): # print("Exists:", os.path.exists(paths[0])) # -------- Settings -------------- ref_img = tio.ScalarImage(paths[0]) ref_img = tio.ToCanonical()(ref_img) # Spacing target_spacing = (0.7, 0.7, 3) ref_img = tio.Resample(target_spacing)(ref_img) # Crop target_shape = (512, 512, 32) padding_constant = ref_img.data.min().item() # Ugly workaround: padding_mode='minimum' calculates the minimum per axis, not globally transform = tio.Compose([ tio.Resample(ref_img), # Resample to reference image to ensure that origin, direction, etc, fit tio.CropOrPad(target_shape, padding_mode=padding_constant), ]) crop_height = crop_breast_height(transform(ref_img)) split_side = { 'right': tio.Crop((256, 0, 0, 0, 0, 0)), 'left': tio.Crop((0, 256, 0, 0, 0, 0))} left_slices = [] right_slices = [] for n, path_img in enumerate(paths): # Read image img = tio.ScalarImage(path_img) # Preprocess (eg. Crop/Pad) padding_constant = img.data.min().item() transform = tio.Compose([ tio.Resample(ref_img), tio.CropOrPad(target_shape, padding_mode=padding_constant), ]) img = transform(img) # Crop bottom and top so that height is 256 and breast is preserved img = crop_height(img) # Split left and right side for side in ['left', 'right']: img_side = split_side[side](img) # img_side.save('left.nii.gz') arr_side = img_side.data.numpy()[0] # [C, D, H, W] → [D, H, W] if side == 'left': left_slices.append(arr_side) else: right_slices.append(arr_side) # Now you have: left_array = np.stack(left_slices, axis = -1) right_array = np.stack(right_slices, axis = -1) return create_relative_array(left_array), create_relative_array(right_array) def create_relative_array(array): base = array[:, :, :, 0] S_early = compute_subtraction(array[:, :, :, 1], base) S_mid = compute_subtraction(array[:, :, :, 2], base) S_late = compute_subtraction(array[:, :, :, 3], base) # Calculate wash-in and wash-out# Assuming S_early, S_mid, S_late have shape (X, Y, Z) enh_stack = np.stack([S_early, S_mid, S_late], axis=-1) # shape: (X, Y, Z, 3) # Compute voxel-wise max across time max_enhancement = np.max(enh_stack, axis=-1) output = np.stack([ base, # Image at index 1 array[:, :, :, 1], # Image at index 1 array[:, :, :, 2], array[:, :, :, 3], S_early, S_mid, S_late, max_enhancement, ], axis=-1) return output