Compress cat12 outputs

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

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)

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"),