diff --git a/brainprep/interfaces/cat12.py b/brainprep/interfaces/cat12.py index cd155095..4b72e356 100644 --- a/brainprep/interfaces/cat12.py +++ b/brainprep/interfaces/cat12.py @@ -12,6 +12,7 @@ """ import glob +import os from pathlib import Path import pandas as pd @@ -36,6 +37,7 @@ ) from ..utils import ( coerce_to_path, + compress_nii, parse_bids_keys, ) from .utils import ( @@ -404,3 +406,73 @@ def cat12vbm_morphometry( morphometry_files.append(volume_file) return (morphometry_files, ) + + +def find_cat12_input_copy( + matlabbatch: File, + ) -> list[File]: + """ + When running, CAT12 copies the image(s) it works on from rawdata as + uncompressed nifti next to the matlabbatch. + The goal of this function is to get the path of these copies from the + matlabbatch content. + + Parameters + ---------- + matlabbatch: File + Path to the ready for execution CAT12 batch file + + Returns + ------- + input_copies: list[File] + Files copied from rawdata and processed by cat12. + """ + if isinstance(matlabbatch, str): + matlabbatch = Path(matlabbatch) + cat12_outdir = matlabbatch.parent + input_copies = [] + with open(matlabbatch) as file: + for line in file: + if str(cat12_outdir) in line and ".nii" in line: + line = line.strip().strip("\'") + input_copies.append(Path(line.strip())) + return input_copies + + +def clean_cat12_outputs( + matlabbatch: File, + gm_files: list[File], + ) -> list[File]: + """ + Compress the nifti outputs produced by cat12 using nibabel, + and remove the rawdata copy next to the matlabbatch file. + + Parameters + ---------- + matlabbatch: File + Path to the ready for execution CAT12 batch file. + gm_files: list[File] + Path to the modulated, normalized gray matter segmentations + outputed by cat12. + + Returns + ------- + gm_files: list[File] + Updated paths to the mwp1 files (with .gz if they have been + compressed). + """ + niis = matlabbatch.parent.glob("mri/*.nii") + for nii in niis: + # compress + nii_gz = compress_nii(str(nii)) + # replace by compressed image path + if nii in gm_files: + gm_files = [nii_gz if path == nii else path + for path in gm_files] + # remove the rawdata copy next to the matlabbatch file + input_t1_copies = find_cat12_input_copy(matlabbatch) + print(f"Found inputs (to delete): {input_t1_copies}") + for input_file in input_t1_copies: + os.remove(input_file) + print(gm_files) + return gm_files diff --git a/brainprep/utils/utils.py b/brainprep/utils/utils.py index f8aea780..a6aa0213 100644 --- a/brainprep/utils/utils.py +++ b/brainprep/utils/utils.py @@ -11,6 +11,7 @@ """ import inspect +import os import re import uuid from pathlib import Path @@ -21,6 +22,9 @@ get_origin, ) +import nibabel as nib +import numpy as np + from ..config import ( DEFAULT_OPTIONS, brainprep_options, @@ -566,3 +570,79 @@ def find_first_occurrence( raise ValueError( f"Unable to find target '{target}' in parents of {input_file}" ) + + +def compress_nii( + mri_nii: File, + atol: float = 1e-2, + clean_nii: bool = True, + keep: str = 'original' + ) -> File | list[File]: + """Compress a .nii to a .nii.gz using nibabel. Before deleting the .nii, + the function checks that the compressed version is not too different. If it + is, keep images follwing the 'keep' keyword. + + Arguments + --------- + mri_nii: str + Image to compress. Must be a .nii. + atol: float + Absolute tolerance for content comparison between the .nii + and the generated .nii.gz + clean_nii: bool, default True + Set to True if you want to delete the .nii after they have + been successfully compressed to .nii.gz. + keep_both: str, default "both" + Set up the behavior when the .nii and the .nii.gz don't contain + the same data. + Set to "both" to keep both the .nii and the .nii.gz. + Set to "orignal" to keep the .nii. + Set to "compressed" to keep the .nii.gz. + """ + # save a copy as nii.gz + if isinstance(mri_nii, str): + mri_niigz = mri_nii.replace('.nii', '.nii.gz') + elif isinstance(mri_nii, Path): + mri_niigz = mri_nii.with_suffix(".gz") + else: + raise TypeError("'mri_nii' should be either str of Path, " + f"not {type(mri_nii)}") + img_nii = nib.load(mri_nii) + nib.save(img_nii, mri_niigz) + # check the copy is good + img_niigz = nib.load(mri_niigz) + equal_affine = (img_nii.affine == img_niigz.affine).all() + similar_content = np.allclose(img_nii.get_fdata(), + img_niigz.get_fdata(), + atol=atol, + equal_nan=True) + # if copy good, delete nii + if equal_affine and similar_content: + if clean_nii: + os.remove(mri_nii) + return mri_niigz + else: + # if copy not good, print a warning and keep files according + # to the keep argument + if not equal_affine: + msg = f"Different affines for {mri_nii}." + elif not similar_content: + adiff = np.abs(img_nii.get_fdata() - img_niigz.get_fdata()).max() + msg = f"Images content too different for {mri_nii}: {adiff}." + else: + raise AssertionError("This shouldn't be triggered.") + msg += f" Keep {keep}" + print(msg) + if keep == "original": + os.remove(mri_niigz) + return mri_nii + elif keep == "compressed": + os.remove(mri_nii) + return mri_niigz + elif keep == "both": + return mri_nii, mri_niigz + else: + err_msg = (f"{keep} value not handled for 'keep' argument. " + "Please choose between 'original', 'compressed' " + "and 'both'.") + raise ValueError(err_msg) diff --git a/brainprep/workflow/vbm.py b/brainprep/workflow/vbm.py index b03624ae..902e0b42 100644 --- a/brainprep/workflow/vbm.py +++ b/brainprep/workflow/vbm.py @@ -116,6 +116,11 @@ def brainprep_vbm( [entities], ) + gm_files = interfaces.clean_cat12_outputs( + batch_file, + gm_files, + ) + interfaces.anonfile( batch_file, find_first_occurrence(output_dir, "derivatives"), @@ -223,6 +228,11 @@ def brainprep_longitudinal_vbm( entities, ) + gm_files = interfaces.clean_cat12_outputs( + batch_file, + gm_files, + ) + interfaces.anonfile( batch_file, find_first_occurrence(output_dir, "derivatives"),