<jupyter_start><jupyter_text>Data preprocessing Data preprocessing for MUDI data and Human Connectome Project (HCP). The goal of this notebook is to create a HDF5 file with the following structure: - `data`: contains all normalized voxel data and attributes to "de'-normalize the data - `index`: a list of indexes to connect each voxel to a subject - `masks`: a list of integers to connect each voxel to a tissue type<jupyter_code>import os
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional

import h5py
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

import torch
from nilearn import image, plotting
from nilearn.masking import apply_mask

from autoencoder.logger import logger, set_log_level<jupyter_output><empty_output><jupyter_text>**Change these to your correct path**Make sure you have the following packages installed: - FSL: https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/<jupyter_code>MUDI_PATH = Path("/media/maarten/disk1/MUDI/")
HCP_PATH = Path("/media/maarten/disk1/HCP_2")

MUDI_OUTPUT_PATH = Path("/media/maarten/disk1/MUDI")
HCP_OUTPUT_PATH = Path("/media/maarten/disk1/HCP_2")

FSL_INSTALL_PATH = "/usr/local/fsl"
os.environ["FSLOUTPUTTYPE"] = "NIFTI_GZ"
os.environ["FSLDIR"] = FSL_INSTALL_PATH
os.environ["PATH"] = f"{FSL_INSTALL_PATH}/bin:" + os.environ["PATH"]
set_log_level(10)<jupyter_output><empty_output><jupyter_text>Step 1: Retrieving file pathsWe create a `dataclass` to store all the paths in. We can use these later to load the files.<jupyter_code>@dataclass
class MRISubjectData:
    subject_name: str
    subject_id: int
    project: str  # Name of the dataset
    output_file: Path

    # paths
    scheme_path: Path
    root_path: Path  # path were all MRI data of the subject is located
    dmri_path: Path  # relative to root_path
    t1_path: Optional[Path] = None  # relative to root_path
    t2_path: Optional[Path] = None  # relative to root_path
    brain_mask_path: Optional[Path] = None  # relative to root_path
    fivett_mask_path: Optional[Path] = None

    # normalization data
    max_data: Optional[float] = 1.0
    lstq_coefficient: Optional[float] = 1.0
mri_mudi_subjects: List[MRISubjectData] = list()
mri_hcp_subjects: List[MRISubjectData] = list()

# MUDI subjects
for i in range(1, 6):
    subject = f"cdmri001{i}"

    # check if we have a previously generated 5tt mask
    fivett_mask_path = None
    if Path(MUDI_PATH, subject, "5tt.nii").exists():
        fivett_mask_path = Path(MUDI_PATH, subject, "5tt.nii")

    mri_mudi_subjects.append(
        MRISubjectData(
            subject,
            10 + i,
            "MUDI",
            Path(MUDI_OUTPUT_PATH, f"prj_MUDI_data.hdf5"),
            Path(MUDI_PATH, "parameters_new.txt"),
            Path(MUDI_PATH, subject),
            Path("MB_Re_t_moco_registered_applytopup.nii.gz"),
            brain_mask_path=Path("brain_mask.nii.gz"),
            # fivett_mask_path=fivett_mask_path,
        )
    )

logger.info(f"found {len(mri_mudi_subjects)} MUDI subjects")

# HCP subjects
for subject_path in HCP_PATH.iterdir():
    if not subject_path.is_dir():
        continue

    subject = subject_path.name

    if int(subject) in (
        286347,
        995174,
        884064,
        186040,
        827052,
        392750,
        569965,
        114116,
        362034,
        701535,
        153934,
        578057,
        154330,
        804646,
        888678,
        180533,
        578158,
        196851,
        206727,
    ):  # remove subjects with few bvecs
        continue

    # check if we have a previously generated 5tt mask
    fivett_mask_path = None
    if Path(HCP_PATH, subject, "5tt.nii").exists():
        fivett_mask_path = Path(HCP_PATH, subject, "5tt.nii")

    # load and combine all MRI parameters
    scheme_path = Path(HCP_PATH, subject, f"T1w/Diffusion/scheme.txt")
    if not scheme_path.exists():
        bvals_path = Path(HCP_PATH, subject, "T1w/Diffusion/bvals")
        bvecs_path = Path(HCP_PATH, subject, "T1w/Diffusion/bvecs")
        bvals = np.loadtxt(bvals_path)[:, np.newaxis]
        bvals = np.around(bvals, -3)
        bvecs = np.loadtxt(bvecs_path).T
        scheme = np.c_[bvecs, bvals]
        np.savetxt(scheme_path, scheme, fmt="%.6f %.6f %.6f %d")

    mri_hcp_subjects.append(
        MRISubjectData(
            subject,
            int(subject),
            "HCP",
            Path(HCP_OUTPUT_PATH, f"prj_HCP_data.hdf5"),
            scheme_path,
            Path(HCP_PATH, subject),
            Path("T1w/Diffusion/data.nii.gz"),
            Path("T1w/T1w_acpc_dc_restore_1.25.nii.gz"),
            brain_mask_path=Path("T1w/Diffusion/nodif_brain_mask.nii.gz"),
            fivett_mask_path=fivett_mask_path,
        )
    )

logger.info(f"found {len(mri_hcp_subjects)} HCP subjects")<jupyter_output><empty_output><jupyter_text>HCP has a lot of subject and training time might take too long if you want to use them all. This function select `n` random samples, evenly distributed between genders. Make sure `n` is an even number.**Note**: it requires that you download the "Behavioral Data" CSV from the "WU-Minn HCP Data - 1200 Subjects" dataset (https://db.humanconnectome.org). Rename it "prj_HCP_metadata.csv" and place it in the HCP download folder.<jupyter_code>def get_HCP_samples(n: int) -> List[int]:
    df = pd.read_csv(Path(HCP_PATH, "prj_HCP_metadata.csv"), index_col="Subject")

    subject_ids = [s.subject_id for s in mri_hcp_subjects]
    subject_dict = {s.subject_id: s for s in mri_hcp_subjects}

    df = df.loc[subject_ids]
    for selected_subject_id in list(df.groupby(by="Gender").sample(n=n // 2).index):
        yield subject_dict[selected_subject_id]


mri_hcp_subjects = list(get_HCP_samples(6))<jupyter_output><empty_output><jupyter_text>Step 2: Creating and applying a brain maskCreate a brain mask from dwi image if none was provided. Based on the script bet from FSL toolbox. Docs: https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/BET/UserGuide<jupyter_code>def generate_brain_mask(subject: MRISubjectData):
    """"""
    logger.info("starting brain mask generation")
    result = subprocess.run(
        [
            "bet",
            Path(subject.root_path, subject.dmri_path),
            Path(subject.root_path, "brain.nii"),
            "-m",
            "-n",
            "-f",
            "0.5",
        ]
    )
    if result.returncode == 0:
        subject.brain_mask_path = Path(subject.root_path, "brain_mask.nii.gz")
    logger.info("finished brain mask generation, return code: %d", result.returncode)
    return result.returncode<jupyter_output><empty_output><jupyter_text>Apply the brain mask.<jupyter_code>def load_scan(subject: MRISubjectData) -> np.ndarray:
    dmri_path = Path(subject.root_path, subject.dmri_path)
    brain_mask_path = None
    if subject.brain_mask_path:
        brain_mask_path = Path(subject.root_path, subject.brain_mask_path)
    else:
        result_code = generate_brain_mask(subject)
        if result_code != 0:
            logger.error("could not generate brain mask")
            return
        brain_mask_path = Path(subject.root_path, subject.brain_mask_path)
    return np.transpose(apply_mask(imgs=str(dmri_path), mask_img=str(brain_mask_path)))
def store_dmri_data(scan_data: np.ndarray, subject: MRISubjectData):
    logger.info("storing data for subject %s", subject.subject_name)

    indexes = np.full((scan_data.shape[0],), subject.subject_id)

    with h5py.File(subject.output_file, "a", libver="latest") as archive:
        if "data" in archive.keys():
            archive["index"].resize((archive["index"].shape[0] + indexes.shape[0]), axis=0)
            archive["index"][-indexes.shape[0] :] = indexes

            archive["data"].resize((archive["data"].shape[0] + scan_data.shape[0]), axis=0)
            archive["data"][-scan_data.shape[0] :] = scan_data
            archive["data"].attrs["max_data"] = np.append(archive["data"].attrs["max_data"], [subject.max_data])
            archive["data"].attrs["lstsq_coefficient"] = np.append(
                archive["data"].attrs["lstsq_coefficient"], [subject.lstsq_coefficient]
            )
            archive["data"].attrs[str(subject.subject_id)] = archive["data"].attrs["max_data"].shape[0] - 1
        else:
            archive.create_dataset("index", data=indexes, chunks=(1024,), maxshape=(None,))
            dataset = archive.create_dataset(
                "data", data=scan_data, chunks=(1024, scan_data.shape[1]), maxshape=(None, scan_data.shape[1])
            )
            dataset.attrs["max_data"] = np.array([subject.max_data])
            dataset.attrs["lstsq_coefficient"] = np.array([subject.lstsq_coefficient])
            dataset.attrs[str(subject.subject_id)] = 0


def normalize_data(subjects: List[MRISubjectData]):
    scheme = np.loadtxt(subjects[0].scheme_path)
    mask_b_0 = scheme[:, 3] == 0.0
    subject_0 = subjects[0]

    logger.info("normalizing subject %s (subject 0)", subject_0.subject_name)

    scan_data_subject_0 = load_scan(subject_0)

    max_data = np.percentile(scan_data_subject_0, 95)
    logger.info("95th percentile: %f", max_data)
    subject_0.max_data = max_data
    subject_0.lstsq_coefficient = 1.0

    scan_data_median_subject_0 = np.median(scan_data_subject_0[:, mask_b_0], axis=0)
    scan_data_normalized = scan_data_subject_0.astype("float32") / subject_0.max_data

    store_dmri_data(scan_data_normalized, subject_0)

    for subject in subjects[1:]:
        logger.info("normalizing subject %s", subject.subject_name)

        # if subject.output_file.exists():
        #     continue

        scan_data = load_scan(subject)
        median_scan = np.median(scan_data[:, mask_b_0], axis=0)

        scan_lstsq_coef, _, _, _ = np.linalg.lstsq(median_scan[:, np.newaxis], scan_data_median_subject_0, rcond=-1)
        logger.info("lstsq coefficient for subject %s: %f", subject.subject_name, scan_lstsq_coef[0])

        subject.lstsq_coefficient = scan_lstsq_coef[0]
        subject.max_data = max_data

        scan_data_normalized = scan_data.astype("float32") * subject.lstsq_coefficient / subject.max_data

        store_dmri_data(scan_data_normalized, subject)
normalize_data(mri_mudi_subjects)
normalize_data(mri_hcp_subjects)<jupyter_output><empty_output><jupyter_text>Step 3: Process bvecs and bvals and create scheme fileAverage b vectors and round b values to the nearest 1000.<jupyter_code>def process_bvecs_bvals(subjects: List[MRISubjectData]):
    bvecs = None
    bvals = None
    other = None

    for subject in subjects:
        scheme = np.loadtxt(subject.scheme_path)
        if bvecs is not None:
            bvecs += scheme[:, :3]
        else:
            bvecs = scheme[:, :3]
            bvals = scheme[:, 3]
            other = scheme[:, 4:]

    bvals = np.around(bvals, -3)
    bvecs /= len(subjects)

    for i in range(bvecs.shape[0]):
        if bvals[i] == 0.0:
            bvecs[i] = 0.0
        else:
            bvecs[i] /= np.sqrt(np.sum(bvecs[i] ** 2))

    return np.c_[bvecs, bvals, other]
scheme = process_bvecs_bvals(mri_mudi_subjects)
np.savetxt(Path(MUDI_PATH, "scheme.txt"), scheme, fmt="%.6f %.6f %.6f %d %d %d")
with h5py.File(Path(MUDI_OUTPUT_PATH, "prj_MUDI_parameters.hdf5"), "w") as hdf5_f:
    hdf5_f.create_dataset("parameters", data=scheme)
scheme = process_bvecs_bvals(mri_hcp_subjects)
np.savetxt(Path(HCP_PATH, "parameters.txt"), scheme, fmt="%.6f %.6f %.6f %d")
with h5py.File(Path(HCP_OUTPUT_PATH, "prj_HCP_parameters.hdf5"), "w") as hdf5_f:
    hdf5_f.create_dataset("parameters", data=scheme)<jupyter_output><empty_output><jupyter_text>Step 4: Generate White matter, grey matter, etc, masks<jupyter_code>def generate_5tt_mask(subject: MRISubjectData):
    """Docs: https://mrtrix.readthedocs.io/en/latest/reference/commands/5ttgen.html#ttgen-fsl"""
    logger.info("starting 5tt generation, this will take a few minutes...")
    output_path = Path(subject.root_path, "5tt.nii")
    result = subprocess.run(
        [
            "5ttgen",
            "fsl",
            Path(subject.root_path, subject.t1_path),  # The input T1-weighted image
            output_path,  # The output 5TT image
            "-mask",
            Path(subject.root_path, subject.brain_mask_path),
            "-force",
            "-nocrop",
        ]
    )
    if result.returncode != 0:
        logger.error("could not finish generation, return code: %d", result.returncode)
        return

    subject.fivett_mask_path = output_path
    logger.info("finished 5tt generation, return code: %d", result.returncode)
    return result.returncode
def generate_tt_masks_by_dwi(subject: MRISubjectData):
    logger.info("starting 5tt generation, this will take a few minutes...")
    output_path = Path(subject.root_path, "3tt.nii")
    # calculate reponse data
    subprocess.run(
        [
            "dwi2response",
            "dhollander",
            Path(subject.root_path, subject.dmri_path),
            Path(subject.root_path, "wm_response.txt"),
            Path(subject.root_path, "gm_response.txt"),
            Path(subject.root_path, "csf_response.txt"),
            "-grad",
            subject.scheme_path,
            "-force",
        ]
    )
    # calculate FOD data
    subprocess.run(
        [
            "dwi2fod",
            "msmt_csd",
            Path(subject.root_path, subject.dmri_path),
            Path(subject.root_path, "wm_response.txt"),
            Path(subject.root_path, "wmfod.nii"),
            Path(subject.root_path, "gm_response.txt"),
            Path(subject.root_path, "gm.nii"),
            Path(subject.root_path, "csf_response.txt"),
            Path(subject.root_path, "csf.nii"),
            "-grad",
            subject.scheme_path,
            "-mask",
            Path(subject.root_path, subject.brain_mask_path),
            "-force",
        ]
    )
    # remove FOD data
    subprocess.run(
        [
            "mrconvert",
            Path(subject.root_path, "wmfod.nii"),
            Path(subject.root_path, "wm.nii"),
            "-coord",
            "3",
            "0",
            "-axes",
            "0,1,2",
            "-force",
        ]
    )
    # combine all scans
    result = subprocess.run(
        [
            "mrcat",
            Path(subject.root_path, "gm.nii"),
            Path(subject.root_path, "wm.nii"),
            Path(subject.root_path, "csf.nii"),
            output_path,
            "-axis",
            "3",
            "-force",
        ]
    )
    
    subject.fivett_mask_path = output_path
    logger.info("finished 5tt generation, return code: %d", result.returncode)
    return result.returncode
def binarize_and_mask_5tt(subject: MRISubjectData):
    target_img = image.load_img(str(Path(subject.root_path, subject.brain_mask_path)))
    target_affine = target_img.affine
    target_shape = image.get_data(target_img).shape

    source_img = image.load_img(str(subject.fivett_mask_path))
    source_affine = source_img.affine
    source_shape = image.get_data(source_img).shape

    # correct 5tt mask affine if its not the same as the brain mask
    # ------
    if not np.allclose(target_affine, source_affine):
        logger.warning("Affine values not the same, correcting...")
        source_img = image.resample_img(source_img, target_affine=target_affine, target_shape=target_shape)

    # apply brain mask
    # ------
    img = np.transpose(apply_mask(source_img, target_img))

    # binarize the 5tt mask
    # ------
    r = img.argmax(axis=1)  # roll index
    _5tt_mask = np.c_[np.ones(img.shape[0]), np.zeros((img.shape[0], img.shape[1] - 1))]

    # shift ones (1) from first column of _5tt_mask to the correct column with the roll index using this code:
    # https://stackoverflow.com/a/20361561
    rows, column_indices = np.ogrid[: _5tt_mask.shape[0], : _5tt_mask.shape[1]]
    column_indices = column_indices - r[:, np.newaxis]
    return _5tt_mask[rows, column_indices]


def create_and_store_5tt(subjects):
    for subject in subjects:
        logger.info("Proccessing subject %s", subject.subject_name)
        if subject.fivett_mask_path is None and subject.t1_path is not None:
            returncode = generate_5tt_mask(subject)
            if returncode != 0:
                logger.error("Could not generate 5tt masks")
                continue
        elif subject.fivett_mask_path is None and subject.t1_path is None:
            generate_tt_masks_by_dwi(subject)
            

        _5tt_bin_path = Path(subject.root_path, "5tt_bin.npy")
        if not _5tt_bin_path.exists():
            _5tt_data = binarize_and_mask_5tt(subject)
            with open(_5tt_bin_path, "wb") as f:
                np.save(f, _5tt_data.T)

        with open(_5tt_bin_path, "rb") as f:
            data = np.load(f)
        
        # Combine all masks into one:
        if data.shape[0] == 5:
            # - 0 = Grey Matter/gm
            # - 1 = White Matter/wm
            # - 2 = CSF
            gm = data[0] + data[1]  # combine Cortical grey matter and Sub-cortical grey matter
            wm = data[2]
            csf = data[3]
        else:
            gm = data[0]
            wm = data[1]
            csf = data[2]
            
        wm[wm == 1] = 2
        csf[csf == 1] = 3
        masks = gm + wm + csf
        masks -= 1
                        
        with h5py.File(subject.output_file, "a", libver="latest") as archive:
            key = "masks"
            if key in archive.keys():
                archive[key].resize((archive[key].shape[0] + masks.shape[0]), axis=0)
                archive[key][-masks.shape[0] :] = masks
            else:
                dataset = archive.create_dataset(key, data=masks, dtype="i8", chunks=(1024,), maxshape=(None,))
                dataset.attrs["gm"] = np.array([0])
                dataset.attrs["wm"] = np.array([1])
                dataset.attrs["csf"] = np.array([2])
                dataset.attrs["wb"] = np.array([0, 1, 2])
create_and_store_5tt(mri_mudi_subjects)
subj = np.concatenate((subj11, subj12, subj13, subj14, subj15), axis=0)
print(subj.shape)
masked_data = np.concatenate(
    (masked_data11n, masked_data12n, masked_data13n, masked_data14n, masked_data15n),
    axis=0,
)
print(masked_data.shape)
import pandas as pd
df = pd.DataFrame(np.concatenate((subj[:, np.newaxis], masked_data), axis=1))
df
df.to_csv("data.csv")
direc16 = "./data"
masked_data16 = np.transpose(
    apply_mask(
        imgs=os.path.join(direc16, "16_MB_RE_t.nii.gz"),
        mask_img=os.path.join(direc16, "brain_mask-testing1.nii.gz"),
    )
)
masked_data16.shape
direc17 = "./data"
masked_data17 = np.transpose(
    apply_mask(
        imgs=os.path.join(direc17, "17_MB_RE_t.nii.gz"),
        mask_img=os.path.join(direc17, "brain_mask-testing2.nii.gz"),
    )
)
med16 = np.median(masked_data16[:, mask2], axis=0)
med17 = np.median(masked_data17[:, mask2], axis=0)
med11_ = np.median(masked_data11[:, mask3], axis=0)
a16, _, _, _ = np.linalg.lstsq(med16[:, np.newaxis], med11_)
a17, _, _, _ = np.linalg.lstsq(med17[:, np.newaxis], med11_)
print(a16, a17)
fig, axes = plt.subplots(1, 2, sharey=True, figsize=(20, 5))
axes[0].plot(med16, med11_, "yo", med16, med16 * a16, ":k", med16, med16, "-k")
axes[1].plot(med17, med11_, "bo", med17, med17 * a17, ":k", med17, med17, "-k")
for ax in axes:
    ax.set(aspect="equal")
fig.show()
selected_imgs = image.index_img(os.path.join(direc11, img_file), np.array(selecind[[1, 10, 100, 300]]))
for img in image.iter_img(selected_imgs):
    # img is now an in-memory 3D img
    plotting.plot_anat(img, vmin=0, vmax=15)
selected_imgs = image.index_img(os.path.join(direc16, "16_MB_RE_t.nii.gz"), np.array([1, 10, 100, 300]))
for img in image.iter_img(selected_imgs):
    # img is now an in-memory 3D img
    plotting.plot_anat(img, vmin=0, vmax=15)
masked_data16n = masked_data16.astype("float32") * a16 / max_data
masked_data17n = masked_data17.astype("float32") * a17 / max_data
subj16 = 16 * np.ones((masked_data16.shape[0],), dtype=int)
subj17 = 17 * np.ones((masked_data17.shape[0],), dtype=int)
subj = np.concatenate((subj16, subj17), axis=0)
print(subj.shape)
masked_data = np.concatenate((masked_data16n, masked_data17n), axis=0)
print(masked_data.shape)
ind = np.arange(len(subj16) + len(subj17))
ind.shape
df1 = pd.DataFrame(np.concatenate((ind[:, np.newaxis], subj[:, np.newaxis]), axis=1))
df1.to_csv("header_test.csv")
import h5py

h5f = h5py.File("data_test.hdf5", "w")
h5f.create_dataset("data1", data=masked_data)
h5f.close
# max_data = masked_data11.max()
max_data = np.percentile(masked_data11, 95)
masked_data11n = masked_data11.astype("float32") / max_data
masked_data12n = masked_data12.astype("float32") * a12 / max_data
masked_data13n = masked_data13.astype("float32") * a13 / max_data
masked_data14n = masked_data14.astype("float32") * a14 / max_data
masked_data15n = masked_data15.astype("float32") * a15 / max_data
print(max_data)
print(masked_data11.shape)
fig, axes = plt.subplots(1, 5, sharey=True, figsize=(20, 10))
axes[0].hist(masked_data11n.flatten(), range=[0, 1])
axes[1].hist(masked_data12n.flatten(), range=[0, 1])
axes[2].hist(masked_data13n.flatten(), range=[0, 1])
axes[3].hist(masked_data14n.flatten(), range=[0, 1])
axes[4].hist(masked_data15n.flatten(), range=[0, 1])
fig.show()
subj11 = 11 * np.ones((masked_data11.shape[0],), dtype=int)
subj12 = 12 * np.ones((masked_data12.shape[0],), dtype=int)
subj13 = 13 * np.ones((masked_data13.shape[0],), dtype=int)
subj14 = 14 * np.ones((masked_data14.shape[0],), dtype=int)
subj15 = 15 * np.ones((masked_data15.shape[0],), dtype=int)
subj = np.concatenate((subj11, subj12, subj13, subj14, subj15), axis=0)
print(subj.shape)
masked_data = np.concatenate(
    (masked_data11n, masked_data12n, masked_data13n, masked_data14n, masked_data15n),
    axis=0,
)
print(masked_data.shape)
ind = np.arange(len(subj11) + len(subj12) + len(subj13) + len(subj14) + len(subj15))
ind.shape
import pandas as pd
df1 = pd.DataFrame(np.concatenate((ind[:, np.newaxis], subj[:, np.newaxis]), axis=1))
df1.to_csv("header_.csv")
df1
df = pd.DataFrame(masked_data)
df.to_csv("data_.csv")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1291122127084272, 'fraction_numerical': 0.03181613339342046, 'mean_word_length': 3.0897530409141174, 'pattern_counts': {'":': 2, '<': 33, '<?xml version=': 0, '>': 35, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '436821', 'license': 'MIT', 'n_tokens_mistral': 8293, 'n_tokens_neox': 7853, 'n_words': 1541, 'path': 'notebooks/data_preprocessing.ipynb', 'repo_name': 'GewoonMaarten/spherical-dmri-conv', 'size': '22190'}