JeeKay/brain-tumor-segmentation

1import os2import h5py3import numpy as np4import tensorflow as tf5from sklearn.model_selection import train_test_split6 7def load_h5_file_multitask(file_path):8    with h5py.File(file_path, 'r') as f:9        image = f['image'][()]  # (240, 240, 4)10        mask = f['mask'][()]    # (240, 240) OR (240, 240, 3)11 12    # Normalize image13    image = (image - np.mean(image, axis=(0, 1), keepdims=True)) / \14            (np.std(image, axis=(0, 1), keepdims=True) + 1e-6)15 16    # Handle case if mask is already one-hot encoded (3-channel)17    if mask.ndim == 3 and mask.shape[-1] == 3:18        # One-hot format (NCR, ED, ET) → derive binary targets19        ncr = mask[..., 0]20        ed = mask[..., 1]21        et = mask[..., 2]22 23        wt = ((ncr + ed + et) > 0).astype(np.float32)[..., np.newaxis]  # Whole tumor24        tc = ((ncr + et) > 0).astype(np.float32)[..., np.newaxis]       # Tumor core25        et = (et > 0).astype(np.float32)[..., np.newaxis]               # Enhancing tumor26    else:27        # Single-channel label map28        mask = mask.astype(np.uint8)29        wt = (mask > 0).astype(np.float32)[..., np.newaxis]30        tc = np.isin(mask, [1, 4]).astype(np.float32)[..., np.newaxis]31        et = (mask == 4).astype(np.float32)[..., np.newaxis]32 33    return image.astype(np.float32), wt, tc, et34 35 36def _parse_multitask_function(path):37    image, wt, tc, et = tf.numpy_function(38        load_h5_file_multitask, [path],39        [tf.float32, tf.float32, tf.float32, tf.float32]40    )41 42    image.set_shape((240, 240, 4))43    wt.set_shape((240, 240, 1))44    tc.set_shape((240, 240, 1))45    et.set_shape((240, 240, 1))46 47    masks = {48        'wt_head': wt,49        'tc_head': tc,50        'et_head': et51    }52 53    return image, masks54 55 56def get_dataset(file_paths, batch_size=8, shuffle=False, num_workers=4):57    dataset = tf.data.Dataset.from_tensor_slices(file_paths)58    if shuffle:59        dataset = dataset.shuffle(buffer_size=1000)60    dataset = dataset.map(_parse_multitask_function, num_parallel_calls=num_workers)61    dataset = dataset.batch(batch_size)62    dataset = dataset.prefetch(tf.data.AUTOTUNE)63    return dataset64 65def get_train_val_datasets(data_dir, batch_size=8, test_size=0.2, random_state=42):66    all_files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith('.h5')]67    train_files, val_files = train_test_split(all_files, test_size=test_size, random_state=random_state)68 69    train_dataset = get_dataset(train_files, batch_size=batch_size, shuffle=True, num_workers=4)70    val_dataset = get_dataset(val_files, batch_size=batch_size, shuffle=False, num_workers=4)71 72    print(f"Total files: {len(all_files)}")73    print(f"Train files: {len(train_files)}")74    print(f"Val files: {len(val_files)}")75 76    return train_dataset, val_dataset77