prepare_and_preprocess/preprocess.py

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