Neuroimaging-Example/src/functions.py

import nibabel as nib
import glob
from scipy.ndimage import zoom
from tqdm import tqdm
from pathlib import PurePath
import numpy as np
import os


# TODO: Add docstring to functions

def load_Nift2np(path):
    """
    Reads a .nii.gz file and loads it into numpy array
    """
    return np.array(nib.load(path).dataobj)


def resize(img, shape, mode='constant', orig_shape=(240, 240, 155)):
    """
    Wrapper for scipy.ndimage.zoom suited for MRI images.
    """
    assert len(shape) == 3, "Can not have more than 3 dimensions"
    factors = (
        shape[0] / orig_shape[0],
        shape[1] / orig_shape[1],
        shape[2] / orig_shape[2]
    )

    # Resize to the given shape
    return zoom(img, factors, mode=mode)


def process_img(img, out_shape=None):
    """
    Preprocess the image.
    Just an example, you can add more preprocessing steps if you wish to.
    """
    if out_shape is not None:
        img = resize(img, out_shape, mode='constant')

    # Normalize the image
    mean = img.mean()
    std = img.std()
    return (img - mean) / std


def save_img(img, raw_img_path):
    # Create output path
    process_img_path = PurePath(raw_img_path.replace('raw', 'processed'))

    # Makedir if necessary
    if not os.path.isdir(process_img_path.parent):
        os.makedirs(process_img_path.parent)

    # Save img to file
    ni_img = nib.Nifti1Image(img, affine=np.eye(4))
    nib.save(ni_img, process_img_path)


def preprocess_label(img, out_shape=None, mode='nearest'):
    """
    Separates out the 3 labels from the segmentation provided, namely:
    GD-enhancing tumor (ET — label 4), the peritumoral edema (ED — label 2))
    and the necrotic and non-enhancing tumor core (NCR/NET — label 1)
    """
    ncr = img == 1  # Necrotic and Non-Enhancing Tumor (NCR/NET)
    ed = img == 2  # Peritumoral Edema (ED)
    et = img == 4  # GD-enhancing Tumor (ET)

    if out_shape:
        ncr = resize(ncr, out_shape, mode=mode)
        ed = resize(ed, out_shape, mode=mode)
        et = resize(et, out_shape, mode=mode)

    return np.array([ncr, ed, et], dtype=np.uint8)


def modalities_path(base_path):
    t1 = glob.glob(os.path.join(base_path + '/*t1.nii.gz'))
    t2 = glob.glob(os.path.join(base_path + '/*t2.nii.gz'))
    flair = glob.glob(os.path.join(base_path + '/*flair.nii.gz'))
    t1ce = glob.glob(os.path.join(base_path + '/*t1ce.nii.gz'))
    seg = glob.glob(os.path.join(base_path + '/*seg.nii.gz'))
    return t1, t2, flair, t1ce, seg


def load_raw_data(t1, t2, flair, t1ce, num_samples, input_shape):
    data: np.array = np.empty((num_samples,) + input_shape, dtype=np.float32)
    for index, item in enumerate(tqdm(zip(t1, t2, flair, t1ce), desc="Loading Data", total=len(t1))):
        data[index] = np.array(
            [process_img(load_Nift2np(modal_path), out_shape=input_shape[1:]) for modal_path in item])
    return data


def load_raw_labels(seg, num_samples, input_shape, output_channels):
    labels: np.array = np.empty((num_samples, output_channels) + input_shape[1:], dtype=np.float32)
    for index, modal_path in enumerate(tqdm(seg, total=len(seg), desc="Loading Labels")):
        labels[index] = preprocess_label(load_Nift2np(modal_path), out_shape=input_shape[1:])[None, ...]
    return labels

def save_data(t1, t2, flair, t1ce, data):
    for index, item in enumerate(tqdm(zip(t1, t2, flair, t1ce), desc="Saving Processed Data", total=len(t1))):
        for modal, modal_path in enumerate(item):
            save_img(data[index][modal], modal_path)


def save_label(labels, seg, file_type):
    for label_index, modal_path in enumerate(tqdm(seg, total=len(seg), desc="Loading Labels")):
        for c_type_index, c_type in enumerate(("_NCR_NET", "_ED", "_ET")):
            save_img(labels[label_index][c_type_index], modal_path.split(file_type)[0] + c_type + file_type)