caideface 0.3.0__tar.gz → 0.3.2__tar.gz

{caideface-0.3.0/src/caideface.egg-info → caideface-0.3.2}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: caideface
-Version: 0.3.0
+Version: 0.3.2
 Summary: MRI defacing pipeline with skull-stripping and affine registration from cai4cai
 Author-email: Lorena Garcia-Foncillas <lorenagarfon00@gmail.com>
-License: MIT
+License: Apache-2.0
 Project-URL: Homepage, https://github.com/cai4cai/defacing_pipeline
 Project-URL: Repository, https://github.com/cai4cai/defacing_pipeline
 Keywords: MRI,defacing,anonymisation,skull-stripping,neuroimaging,NER,text-anonymization
@@ -13,7 +13,6 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-License-File: LICENSE.md
 Requires-Dist: nibabel>=4.0
 Requires-Dist: numpy<2,>=1.22
 Requires-Dist: scipy>=1.9
@@ -27,7 +26,8 @@ Requires-Dist: faker>=18.0
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
-Dynamic: license-file
+Provides-Extra: ct
+Requires-Dist: TotalSegmentator; extra == "ct"
 
 # caideface
 
@@ -330,4 +330,4 @@ If you use the text anonymisation (NER + HIPS), please also cite:
 
 ## License
 
-This project is licensed under the MIT License -- see the [LICENSE](LICENSE.md) file for details.
+This project is licensed under the Apache License 2.0 -- see the [LICENSE](https://github.com/cai4cai/defacing_pipeline/blob/main/LICENSE) file for details.

{caideface-0.3.0 → caideface-0.3.2}/README.md

@@ -299,4 +299,4 @@ If you use the text anonymisation (NER + HIPS), please also cite:
 
 ## License
 
-This project is licensed under the MIT License -- see the [LICENSE](LICENSE.md) file for details.
+This project is licensed under the Apache License 2.0 -- see the [LICENSE](https://github.com/cai4cai/defacing_pipeline/blob/main/LICENSE) file for details.

{caideface-0.3.0 → caideface-0.3.2}/pyproject.toml

@@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "caideface"
-version = "0.3.0"
+version = "0.3.2"
 description = "MRI defacing pipeline with skull-stripping and affine registration from cai4cai"
 readme = "README.md"
-license = {text = "MIT"}
+license = {text = "Apache-2.0"}
 requires-python = ">=3.9"
 authors = [
     {name = "Lorena Garcia-Foncillas", email = "lorenagarfon00@gmail.com"},
@@ -38,6 +38,9 @@ dev = [
     "pytest",
     "ruff",
 ]
+ct = [
+    "TotalSegmentator",
+]
 
 [project.scripts]
 caideface = "caideface.cli:main"

{caideface-0.3.0 → caideface-0.3.2}/src/caideface/__init__.py

@@ -8,13 +8,14 @@ A three-step pipeline for anonymising head MRI scans:
 Plus standalone text anonymisation via NER + HIPS (Hiding in Plain Sight).
 """
 
-__version__ = "0.3.0"
+__version__ = "0.3.2"
 
 from .pipeline import DefacePipeline
 from .reorient import reorient_batch, reorient_single
 from .skull_strip import skull_strip_batch, skull_strip_single
 from .register import deface_batch, deface_single
 from .anonymize import anonymize_batch, anonymize_single, default_ner_model_path
+from .background import detect_background_value
 
 __all__ = [
     "DefacePipeline",
@@ -27,4 +28,5 @@ __all__ = [
     "anonymize_batch",
     "anonymize_single",
     "default_ner_model_path",
+    "detect_background_value",
 ]

{caideface-0.3.0 → caideface-0.3.2}/src/caideface/cli.py

@@ -37,25 +37,28 @@ def main():
     run_parser = subparsers.add_parser("run", help="Run the full defacing pipeline", parents=[parent])
     run_parser.add_argument("input_dir", help="Directory containing raw NIfTI files")
     run_parser.add_argument("output_dir", help="Root output directory")
+    run_parser.add_argument("--modality", required=True, choices=["mri", "ct"], help="Image modality (required)")
     run_parser.add_argument("--brainsfit", required=True, help="Path to BRAINSFit executable")
     run_parser.add_argument("--brainsresample", required=True, help="Path to BRAINSResample executable")
     run_parser.add_argument("--device", default=None, choices=["cpu", "cuda"], help="Device for HD-BET (auto-detected if omitted)")
     run_parser.add_argument("--no-tta", action="store_true", default=True, help="Disable HD-BET test-time augmentation (default: disabled)")
     run_parser.add_argument("--dilation-mm", type=float, default=14.0, help="Brain mask dilation in mm (default: 14)")
-    run_parser.add_argument("--background", type=float, default=0, help="Background value for defaced voxels (default: 0 for MRI, use -1024 for CT)")
+    run_parser.add_argument("--background", type=float, default=None, help="Background value for defaced voxels (auto-detected per volume if omitted)")
     run_parser.add_argument("--template", default=None, help="Custom MNI152 skull-stripped template (uses bundled if omitted)")
     run_parser.add_argument("--face-mask", default=None, help="Custom face mask in MNI152 space (uses bundled if omitted)")
     run_parser.add_argument("--steps", default="all", help="Steps to run: all, or comma-separated: reorient,skull_strip,deface")
 
     # --- reorient ---
-    reorient_parser = subparsers.add_parser("reorient", help="Step 1: Reorient NIfTI scans to MNI152", parents=[parent])
+    reorient_parser = subparsers.add_parser("reorient", help="Step 1: Reorient NIfTI scans", parents=[parent])
     reorient_parser.add_argument("input_dir", help="Directory with NIfTI files")
     reorient_parser.add_argument("output_dir", help="Output directory for reoriented files")
+    reorient_parser.add_argument("--modality", required=True, choices=["mri", "ct"], help="Image modality (mri→LAS, ct→RAS)")
 
     # --- skull-strip ---
     ss_parser = subparsers.add_parser("skull-strip", help="Step 2: Skull-strip with HD-BET", parents=[parent])
     ss_parser.add_argument("input_dir", help="Directory with reoriented NIfTI files")
     ss_parser.add_argument("output_dir", help="Output directory for HD-BET results")
+    ss_parser.add_argument("--modality", required=True, choices=["mri", "ct"], help="Image modality (required)")
     ss_parser.add_argument("--device", default=None, choices=["cpu", "cuda"], help="Device for HD-BET")
     ss_parser.add_argument("--no-tta", action="store_true", default=True, help="Disable test-time augmentation")
     ss_parser.add_argument("--dilation-mm", type=float, default=14.0, help="Dilation in mm")
@@ -65,11 +68,12 @@ def main():
     deface_parser.add_argument("reoriented_dir", help="Directory with reoriented scans (Step 1 output)")
     deface_parser.add_argument("hdbet_dir", help="Directory with HD-BET results (Step 2 output)")
     deface_parser.add_argument("output_dir", help="Output directory for defaced scans")
+    deface_parser.add_argument("--modality", required=True, choices=["mri", "ct"], help="Image modality (required)")
     deface_parser.add_argument("--brainsfit", required=True, help="Path to BRAINSFit executable")
     deface_parser.add_argument("--brainsresample", required=True, help="Path to BRAINSResample executable")
     deface_parser.add_argument("--template", default=None, help="Custom MNI152 skull-stripped template")
     deface_parser.add_argument("--face-mask", default=None, help="Custom face mask in MNI152 space")
-    deface_parser.add_argument("--background", type=float, default=0, help="Background value (default: 0 for MRI, use -1024 for CT)")
+    deface_parser.add_argument("--background", type=float, default=None, help="Background value (auto-detected per volume if omitted)")
 
     # --- anonymize (batch) ---
     anon_parser = subparsers.add_parser("anonymize", help="Anonymize personal names in all .txt files in a directory", parents=[parent])
@@ -99,6 +103,7 @@ def main():
         pipeline = DefacePipeline(
             brainsfit_path=args.brainsfit,
             brainsresample_path=args.brainsresample,
+            modality=args.modality,
             device=args.device,
             disable_tta=args.no_tta,
             desired_dilation_mm=args.dilation_mm,
@@ -113,12 +118,13 @@ def main():
             sys.exit(1)
 
     elif args.command == "reorient":
-        reorient_batch(args.input_dir, args.output_dir)
+        reorient_batch(args.input_dir, args.output_dir, modality=args.modality)
 
     elif args.command == "skull-strip":
         skull_strip_batch(
             args.input_dir,
             args.output_dir,
+            modality=args.modality,
             device=args.device,
             disable_tta=args.no_tta,
             desired_dilation_mm=args.dilation_mm,
@@ -131,6 +137,7 @@ def main():
             output_dir=args.output_dir,
             brainsfit_path=args.brainsfit,
             brainsresample_path=args.brainsresample,
+            modality=args.modality,
             target_path=args.template,
             face_mask_path=args.face_mask,
             background_value=args.background,

{caideface-0.3.0 → caideface-0.3.2}/src/caideface/pipeline.py

@@ -13,7 +13,7 @@ logger = logging.getLogger(__name__)
 
 
 class DefacePipeline:
-    """Orchestrates the three-step MRI defacing pipeline.
+    """Orchestrates the three-step defacing pipeline.
 
     Parameters
     ----------
@@ -21,14 +21,16 @@ class DefacePipeline:
         Path to the BRAINSFit executable (from 3D Slicer).
     brainsresample_path : str
         Path to the BRAINSResample executable (from 3D Slicer).
+    modality : str
+        Image modality: ``'mri'`` or ``'ct'``.
     device : str or None
         'cpu' or 'cuda' for HD-BET. Auto-detected if None.
     disable_tta : bool
         Disable HD-BET test-time augmentation for faster processing.
     desired_dilation_mm : float
         Physical dilation in mm for brain mask expansion.
-    background_value : float
-        Value for defaced voxels (default 0 for MRI; use -1024 for CT).
+    background_value : float or None
+        Value for defaced voxels. Auto-detected per volume when None.
     target_path : str or None
         Custom MNI152 skull-stripped template. Uses bundled if None.
     face_mask_path : str or None
@@ -39,15 +41,17 @@ class DefacePipeline:
         self,
         brainsfit_path: str,
         brainsresample_path: str,
+        modality: str = "mri",
         device: str | None = None,
         disable_tta: bool = True,
         desired_dilation_mm: float = 14.0,
-        background_value: float = 0,
+        background_value: float | None = None,
         target_path: str | None = None,
         face_mask_path: str | None = None,
     ):
         self.brainsfit_path = brainsfit_path
         self.brainsresample_path = brainsresample_path
+        self.modality = modality
         self.device = device
         self.disable_tta = disable_tta
         self.desired_dilation_mm = desired_dilation_mm
@@ -93,9 +97,10 @@ class DefacePipeline:
         # Step 1: Reorientation
         if "reorient" in run_steps:
             logger.info("=" * 60)
-            logger.info("STEP 1: Reorientation to MNI152")
+            target = "RAS" if self.modality == "ct" else "LAS"
+            logger.info("STEP 1: Reorientation to %s", target)
             logger.info("=" * 60)
-            reorient_log = reorient_batch(input_dir, reoriented_dir)
+            reorient_log = reorient_batch(input_dir, reoriented_dir, modality=self.modality)
             results["reorient_log"] = reorient_log
         else:
             logger.info("Skipping Step 1 (reorientation)")
@@ -103,11 +108,13 @@ class DefacePipeline:
         # Step 2: Skull-stripping
         if "skull_strip" in run_steps:
             logger.info("=" * 60)
-            logger.info("STEP 2: Skull-stripping with HD-BET")
+            backend = "TotalSegmentator" if self.modality == "ct" else "HD-BET"
+            logger.info("STEP 2: Skull-stripping with %s", backend)
             logger.info("=" * 60)
             skull_strip_log = skull_strip_batch(
                 input_dir=reoriented_dir,
                 output_dir=hdbet_dir,
+                modality=self.modality,
                 device=self.device,
                 disable_tta=self.disable_tta,
                 desired_dilation_mm=self.desired_dilation_mm,
@@ -127,6 +134,7 @@ class DefacePipeline:
                 output_dir=defaced_dir,
                 brainsfit_path=self.brainsfit_path,
                 brainsresample_path=self.brainsresample_path,
+                modality=self.modality,
                 target_path=self.target_path,
                 face_mask_path=self.face_mask_path,
                 background_value=self.background_value,

{caideface-0.3.0 → caideface-0.3.2}/src/caideface/register.py

@@ -11,6 +11,8 @@ import numpy as np
 from numpy.linalg import inv
 from natsort import natsorted
 
+from .background import detect_background_value
+
 logger = logging.getLogger(__name__)
 
 
@@ -23,14 +25,35 @@ def _data_dir() -> str:
     return os.path.join(os.path.dirname(__file__), "data")
 
 
-def default_template_path() -> str:
+def default_template_path(modality: str = "mri") -> str:
+    """Return the default template path for the given modality."""
+    if modality == "ct":
+        return default_ct_template_path()
     return os.path.join(_data_dir(), "mni_icbm152_t1_tal_nlin_sym_55_ext_brain_only.nii.gz")
 
 
-def default_face_mask_path() -> str:
+def default_face_mask_path(modality: str = "mri") -> str:
+    """Return the default face mask path for the given modality."""
+    if modality == "ct":
+        return os.path.join(_data_dir(), "ct_face_mask.nii.gz")
     return os.path.join(_data_dir(), "t1_mask.nii.gz")
 
 
+def default_ct_template_path() -> str:
+    """Return the CT brain atlas from the TotalSegmentator installation."""
+    try:
+        import totalsegmentator
+        ts_dir = os.path.dirname(totalsegmentator.__file__)
+        ct_atlas = os.path.join(ts_dir, "resources", "ct_brain_atlas_1mm.nii.gz")
+        if os.path.isfile(ct_atlas):
+            return ct_atlas
+    except ImportError:
+        pass
+    raise FileNotFoundError(
+        "CT brain atlas not found. Install TotalSegmentator with: pip install caideface[ct]"
+    )
+
+
 # ---------------------------------------------------------------------------
 # Transform I/O (supports both plain 4x4 text and ITK/Slicer formats)
 # ---------------------------------------------------------------------------
@@ -113,6 +136,7 @@ def run_brainsfit(
     output_volume_path: str,
     output_transform_path: str,
     brainsfit_path: str,
+    background_value: float = 0,
 ) -> int:
     """Run BRAINSFit affine registration and return the exit code."""
     cmd = (
@@ -129,7 +153,7 @@ def run_brainsfit(
         f'--medianFilterSize 0,0,0 '
         f'--removeIntensityOutliers 0 '
         f'--outputVolumePixelType float '
-        f'--backgroundFillValue 0 '
+        f'--backgroundFillValue {background_value} '
         f'--interpolationMode Linear '
         f'--numberOfIterations 1500 '
         f'--maximumStepLength 0.05 '
@@ -219,8 +243,9 @@ def deface_single(
     face_mask_path: str,
     brainsfit_path: str,
     brainsresample_path: str,
+    modality: str = "mri",
     existing_transform: str | None = None,
-    background_value: float = 0,
+    background_value: float | None = None,
 ) -> bool:
     """Register, warp face mask, and deface a single scan.
 
@@ -236,6 +261,20 @@ def deface_single(
         logger.info("Already defaced, skipping: %s", masked_path)
         return True
 
+    # --- Background value detection (always runs per volume) ---
+    floating_nii = nib.load(reoriented_path)
+    floating_data = floating_nii.get_fdata()
+
+    detected_bg = detect_background_value(floating_data, modality)
+    if background_value is not None:
+        logger.info(
+            "Detected background: %.1f, using override: %.1f",
+            detected_bg, background_value,
+        )
+    else:
+        background_value = detected_bg
+        logger.info("Using detected background: %.1f", background_value)
+
     out_affine = os.path.join(results_dir, base.replace(".nii.gz", ".txt"))
     out_resampled = os.path.join(results_dir, base.replace(".nii.gz", "_resampled.nii.gz"))
 
@@ -246,7 +285,8 @@ def deface_single(
     else:
         logger.info("Running BRAINSFit registration...")
         exit_code = run_brainsfit(
-            target_path, floating_path, out_resampled, out_affine, brainsfit_path
+            target_path, floating_path, out_resampled, out_affine, brainsfit_path,
+            background_value=background_value,
         )
         if exit_code != 0:
             logger.error("BRAINSFit failed (exit %d) for %s", exit_code, floating_path)
@@ -267,9 +307,6 @@ def deface_single(
     )
 
     # --- Combine masks and deface ---
-    floating_nii = nib.load(reoriented_path)
-    floating_data = floating_nii.get_fdata()
-
     face_data = nib.load(face_mask_resampled).get_fdata()
     brain_data = nib.load(brain_mask_path).get_fdata()
 
@@ -291,9 +328,10 @@ def deface_batch(
     output_dir: str,
     brainsfit_path: str,
     brainsresample_path: str,
+    modality: str = "mri",
     target_path: str | None = None,
     face_mask_path: str | None = None,
-    background_value: float = 0,
+    background_value: float | None = None,
 ) -> list[str]:
     """Run the full defacing pipeline (Step 3) on all dilated scans.
 
@@ -309,12 +347,14 @@ def deface_batch(
         Path to the BRAINSFit executable.
     brainsresample_path : str
         Path to the BRAINSResample executable.
+    modality : str
+        Image modality: ``'mri'`` or ``'ct'``.
     target_path : str or None
         Path to skull-stripped MNI152 template. Uses bundled if None.
     face_mask_path : str or None
         Path to face mask in MNI152 space. Uses bundled if None.
-    background_value : float
-        Value to fill defaced regions (default 0 for MRI; use -1024 for CT).
+    background_value : float or None
+        Value to fill defaced regions. Auto-detected per volume when None.
 
     Returns
     -------
@@ -322,9 +362,9 @@ def deface_batch(
         Paths of scans that failed to deface.
     """
     if target_path is None:
-        target_path = default_template_path()
+        target_path = default_template_path(modality)
     if face_mask_path is None:
-        face_mask_path = default_face_mask_path()
+        face_mask_path = default_face_mask_path(modality)
 
     hdbet_dir = os.path.abspath(hdbet_dir)
     reoriented_dir = os.path.abspath(reoriented_dir)
@@ -371,6 +411,7 @@ def deface_batch(
                 face_mask_path=face_mask_path,
                 brainsfit_path=brainsfit_path,
                 brainsresample_path=brainsresample_path,
+                modality=modality,
                 existing_transform=existing_transforms.get(fimg),
                 background_value=background_value,
             )

{caideface-0.3.0 → caideface-0.3.2}/src/caideface/reorient.py

@@ -1,4 +1,4 @@
-"""Step 1: Reorientation of NIfTI scans to LAS (MNI152 standard) using nibabel."""
+"""Step 1: Reorientation of NIfTI scans using nibabel."""
 
 import os
 import logging
@@ -9,15 +9,18 @@ import pandas as pd
 
 logger = logging.getLogger(__name__)
 
-# Target orientation matching fslreorient2std / MNI152 standard
-TARGET_ORIENTATION = ("L", "A", "S")
+# Target orientations per modality
+TARGET_ORIENTATIONS = {
+    "mri": ("L", "A", "S"),  # LAS — matches fslreorient2std
+    "ct": ("R", "A", "S"),   # RAS — matches CT brain atlas
+}
 
 
-def reorient_single(input_file: str, output_file: str) -> bool:
-    """Reorient a single NIfTI file to LAS orientation (MNI152 standard).
+def reorient_single(input_file: str, output_file: str, modality: str = "mri") -> bool:
+    """Reorient a single NIfTI file to the target orientation for *modality*.
 
-    This is equivalent to FSL's ``fslreorient2std`` but implemented purely
-    in Python using nibabel, removing the FSL dependency.
+    MRI is reoriented to LAS (matching ``fslreorient2std`` / MNI152).
+    CT is reoriented to RAS (matching the CT brain atlas).
 
     Parameters
     ----------
@@ -25,6 +28,8 @@ def reorient_single(input_file: str, output_file: str) -> bool:
         Path to the input NIfTI (.nii.gz) file.
     output_file : str
         Path where the reoriented file will be saved.
+    modality : str
+        ``'mri'`` (target LAS) or ``'ct'`` (target RAS).
 
     Returns
     -------
@@ -33,10 +38,12 @@ def reorient_single(input_file: str, output_file: str) -> bool:
     """
     os.makedirs(os.path.dirname(output_file), exist_ok=True)
 
+    target = TARGET_ORIENTATIONS.get(modality, TARGET_ORIENTATIONS["mri"])
+
     try:
         img = nib.load(input_file)
         orig_ornt = io_orientation(img.affine)
-        target_ornt = axcodes2ornt(TARGET_ORIENTATION)
+        target_ornt = axcodes2ornt(target)
         transform = ornt_transform(orig_ornt, target_ornt)
         reoriented = img.as_reoriented(transform)
         nib.save(reoriented, output_file)
@@ -47,7 +54,7 @@ def reorient_single(input_file: str, output_file: str) -> bool:
     return os.path.exists(output_file)
 
 
-def reorient_batch(input_dir: str, output_dir: str) -> pd.DataFrame:
+def reorient_batch(input_dir: str, output_dir: str, modality: str = "mri") -> pd.DataFrame:
     """Reorient all NIfTI files found recursively under *input_dir*.
 
     The directory structure is mirrored under *output_dir*.
@@ -59,6 +66,8 @@ def reorient_batch(input_dir: str, output_dir: str) -> pd.DataFrame:
     output_dir : str
         Root directory where reoriented files will be saved,
         preserving the subdirectory structure.
+    modality : str
+        ``'mri'`` (target LAS) or ``'ct'`` (target RAS).
 
     Returns
     -------
@@ -68,6 +77,9 @@ def reorient_batch(input_dir: str, output_dir: str) -> pd.DataFrame:
     input_dir = os.path.abspath(input_dir)
     output_dir = os.path.abspath(output_dir)
 
+    target = TARGET_ORIENTATIONS.get(modality, TARGET_ORIENTATIONS["mri"])
+    logger.info("Target orientation: %s (%s)", "".join(target), modality)
+
     log_rows = []
     for root, _dirs, files in os.walk(input_dir):
         for fname in files:
@@ -79,7 +91,7 @@ def reorient_batch(input_dir: str, output_dir: str) -> pd.DataFrame:
             out_path = os.path.join(output_dir, rel, fname)
 
             logger.info("Reorienting %s", input_path)
-            success = reorient_single(input_path, out_path)
+            success = reorient_single(input_path, out_path, modality=modality)
             log_rows.append({"input": input_path, "output": out_path, "success": success})
 
             if success:

{caideface-0.3.0 → caideface-0.3.2}/src/caideface/skull_strip.py

@@ -82,15 +82,81 @@ def is_valid_3d_volume(filepath: str) -> bool:
         return False
 
 
+# ---------------------------------------------------------------------------
+# Brain extraction backends
+# ---------------------------------------------------------------------------
+
+def _extract_brain_hdbet(
+    input_file: str, output_file: str, device: str = "cpu", disable_tta: bool = True
+) -> None:
+    """Run HD-BET brain extraction (MRI)."""
+    hdbet_bin = _hd_bet_executable()
+    cmd = [hdbet_bin, "-i", input_file, "-o", output_file, "-device", device]
+    if disable_tta:
+        cmd.append("--disable_tta")
+    subprocess.run(cmd, check=True, capture_output=True, text=True)
+
+
+def _extract_brain_totalseg(
+    input_file: str, output_file: str, binary_mask_file: str, device: str = "cpu"
+) -> None:
+    """Run TotalSegmentator brain extraction (CT).
+
+    Saves both a brain-extracted volume (*output_file*) and the raw binary
+    brain mask (*binary_mask_file*).  The mask is needed because CT brain-
+    extracted volumes have negative values, so ``data > 0`` cannot be used
+    to recover the mask (unlike MRI).
+
+    Requires ``pip install caideface[ct]``.
+    """
+    try:
+        from totalsegmentator.python_api import totalsegmentator
+    except ImportError:
+        raise ImportError(
+            "TotalSegmentator is required for CT skull-stripping.\n"
+            "Install it with: pip install caideface[ct]"
+        )
+
+    ts_device = "cpu" if device == "cpu" else "gpu"
+
+    input_img = nib.load(input_file)
+    brain_mask = totalsegmentator(
+        input=input_img,
+        task="total",
+        roi_subset=["brain"],
+        device=ts_device,
+    )
+
+    # Save the raw binary brain mask
+    mask_data = brain_mask.get_fdata().astype(np.uint8)
+    nib.save(
+        nib.Nifti1Image(mask_data, input_img.affine, input_img.header),
+        binary_mask_file,
+    )
+
+    # Apply brain mask to get brain-extracted volume
+    data = input_img.get_fdata()
+    brain_data = data * mask_data
+    nib.save(
+        nib.Nifti1Image(brain_data, input_img.affine, input_img.header),
+        output_file,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Single / batch skull-stripping
+# ---------------------------------------------------------------------------
+
 def skull_strip_single(
     input_file: str,
     input_root: str,
     output_dir: str,
+    modality: str = "mri",
     device: str = "cpu",
     disable_tta: bool = True,
     desired_dilation_mm: float = 14.0,
 ) -> dict:
-    """Run HD-BET on a single NIfTI file and produce dilated brain mask.
+    """Extract brain, create dilated mask, and apply it to a single NIfTI file.
 
     Parameters
     ----------
@@ -100,58 +166,71 @@ def skull_strip_single(
         Root input directory (used to compute relative paths).
     output_dir : str
         Root output directory.
+    modality : str
+        ``'mri'`` (uses HD-BET) or ``'ct'`` (uses TotalSegmentator).
     device : str
         'cpu' or 'cuda'.
     disable_tta : bool
-        Disable test-time augmentation for faster processing.
+        Disable HD-BET test-time augmentation (MRI only).
     desired_dilation_mm : float
         Physical dilation size in mm for mask expansion.
 
     Returns
     -------
     dict
-        Status dict with keys: input, output, hd_bet, mask, dilated.
+        Status dict with keys: input, output, brain_extract, mask, dilated.
     """
     filename = os.path.basename(input_file)
     stem = filename.replace(".nii.gz", "")
     subfolder = os.path.relpath(os.path.dirname(input_file), start=input_root)
 
-    hd_bet_file = os.path.join(output_dir, subfolder, f"{stem}_brain.nii.gz")
+    brain_file = os.path.join(output_dir, subfolder, f"{stem}_brain.nii.gz")
     mask_file = os.path.join(output_dir, subfolder, f"{stem}_mask.nii.gz")
     dilated_file = os.path.join(output_dir, subfolder, f"{stem}_dilated.nii.gz")
 
-    os.makedirs(os.path.dirname(hd_bet_file), exist_ok=True)
+    os.makedirs(os.path.dirname(brain_file), exist_ok=True)
 
     status = {
-        "hd_bet": "skipped" if os.path.exists(hd_bet_file) else "pending",
+        "brain_extract": "skipped" if os.path.exists(brain_file) else "pending",
         "mask": "skipped" if os.path.exists(mask_file) else "pending",
         "dilated": "skipped" if os.path.exists(dilated_file) else "pending",
     }
 
     try:
-        # --- HD-BET ---
-        if not os.path.exists(hd_bet_file):
-            hdbet_bin = _hd_bet_executable()
-            cmd = [hdbet_bin, "-i", input_file, "-o", hd_bet_file, "-device", device]
-            if disable_tta:
-                cmd.append("--disable_tta")
-            try:
-                subprocess.run(cmd, check=True, capture_output=True, text=True)
-                status["hd_bet"] = "success"
-            except subprocess.CalledProcessError as e:
-                logger.error("hd-bet failed for %s: %s", input_file, e.stderr)
-                status["hd_bet"] = "failed"
-                return {"input": input_file, "output": hd_bet_file, **status, "error": e.stderr}
+        # --- Brain extraction (modality-dependent) ---
+        # For CT, a separate binary mask file is saved alongside the
+        # brain-extracted volume because CT values include negatives,
+        # making ``data > 0`` unreliable for mask recovery.
+        ct_binary_mask_file = os.path.join(
+            output_dir, subfolder, f"{stem}_brain_binary_mask.nii.gz"
+        )
+        if not os.path.exists(brain_file):
+            if modality == "ct":
+                logger.info("Running TotalSegmentator brain extraction...")
+                _extract_brain_totalseg(input_file, brain_file, ct_binary_mask_file, device)
+            else:
+                logger.info("Running HD-BET brain extraction...")
+                try:
+                    _extract_brain_hdbet(input_file, brain_file, device, disable_tta)
+                except subprocess.CalledProcessError as e:
+                    logger.error("HD-BET failed for %s: %s", input_file, e.stderr)
+                    status["brain_extract"] = "failed"
+                    return {"input": input_file, "output": brain_file, **status, "error": e.stderr}
+            status["brain_extract"] = "success"
 
         # --- Binary mask + dilation ---
         if not os.path.exists(mask_file):
-            img = nib.load(hd_bet_file)
-            data = img.get_fdata()
-            voxel_sizes = img.header.get_zooms()[:3]
+            if modality == "ct" and os.path.exists(ct_binary_mask_file):
+                # Use the raw TotalSegmentator binary mask directly
+                img = nib.load(ct_binary_mask_file)
+                binary_volume = (img.get_fdata() > 0).astype(np.uint8)
+            else:
+                # MRI: threshold the brain-extracted volume
+                img = nib.load(brain_file)
+                binary_volume = (img.get_fdata() > 0).astype(np.uint8)
 
+            voxel_sizes = img.header.get_zooms()[:3]
             struct_elem = get_safe_structuring_element(voxel_sizes, desired_dilation_mm)
-
-            binary_volume = (data > 0).astype(np.uint8)
             dilated_data = binary_dilation(binary_volume, structure=struct_elem)
 
             nib.save(nib.Nifti1Image(dilated_data, img.affine, img.header), mask_file)
@@ -165,7 +244,12 @@ def skull_strip_single(
             original_img = nib.load(input_file)
             original_data = original_img.get_fdata()
 
-            dilated_volume = original_data * mask_data
+            # Detect background so outside-mask voxels keep the correct
+            # intensity (e.g. ~-1000 HU for CT instead of 0).
+            from .background import detect_background_value
+            bg = detect_background_value(original_data, modality)
+
+            dilated_volume = np.where(mask_data, original_data, bg)
             nib.save(
                 nib.Nifti1Image(dilated_volume, original_img.affine, original_img.header),
                 dilated_file,
@@ -178,18 +262,20 @@ def skull_strip_single(
             if status[k] == "pending":
                 status[k] = "failed"
 
-    return {"input": input_file, "output": hd_bet_file, **status}
+    return {"input": input_file, "output": brain_file, **status}
 
 
 def skull_strip_batch(
     input_dir: str,
     output_dir: str,
+    modality: str = "mri",
     device: str | None = None,
     disable_tta: bool = True,
     desired_dilation_mm: float = 14.0,
 ) -> pd.DataFrame:
-    """Run HD-BET skull-stripping on all 3D NIfTI volumes under *input_dir*.
+    """Run skull-stripping on all 3D NIfTI volumes under *input_dir*.
 
+    Uses HD-BET for MRI or TotalSegmentator for CT, based on *modality*.
     Scans that are not 3D volumes (2D slices or 4D time series) are skipped.
 
     Parameters
@@ -197,11 +283,13 @@ def skull_strip_batch(
     input_dir : str
         Root directory with reoriented NIfTI files.
     output_dir : str
-        Root output directory for HD-BET results.
+        Root output directory for skull-stripping results.
+    modality : str
+        ``'mri'`` (HD-BET) or ``'ct'`` (TotalSegmentator).
     device : str or None
         'cpu' or 'cuda'. Auto-detected if None.
     disable_tta : bool
-        Disable test-time augmentation.
+        Disable test-time augmentation (MRI/HD-BET only).
     desired_dilation_mm : float
         Physical dilation in mm.
 
@@ -210,8 +298,17 @@ def skull_strip_batch(
     pd.DataFrame
         Processing log.
     """
-    # Validate hd-bet is available (will raise if not found)
-    _hd_bet_executable()
+    # Validate backend is available
+    if modality == "mri":
+        _hd_bet_executable()
+    elif modality == "ct":
+        try:
+            import totalsegmentator  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "TotalSegmentator is required for CT skull-stripping.\n"
+                "Install it with: pip install caideface[ct]"
+            )
 
     if device is None:
         device = get_default_device()
@@ -231,18 +328,22 @@ def skull_strip_batch(
 
             logger.info("Processing: %s", input_file)
             result = skull_strip_single(
-                input_file, input_dir, output_dir, device, disable_tta, desired_dilation_mm
+                input_file, input_dir, output_dir,
+                modality=modality,
+                device=device,
+                disable_tta=disable_tta,
+                desired_dilation_mm=desired_dilation_mm,
             )
             log_data.append(result)
 
     df = pd.DataFrame(log_data)
-    log_path = os.path.join(output_dir, "hd_bet_log.csv")
+    log_path = os.path.join(output_dir, "skull_strip_log.csv")
 
     if os.path.exists(log_path):
         existing = pd.read_csv(log_path)
         df = pd.concat([existing, df], ignore_index=True)
 
     df.to_csv(log_path, index=False)
-    logger.info("HD-BET log saved at %s", log_path)
+    logger.info("Skull-stripping log saved at %s", log_path)
 
     return df

{caideface-0.3.0 → caideface-0.3.2/src/caideface.egg-info}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: caideface
-Version: 0.3.0
+Version: 0.3.2
 Summary: MRI defacing pipeline with skull-stripping and affine registration from cai4cai
 Author-email: Lorena Garcia-Foncillas <lorenagarfon00@gmail.com>
-License: MIT
+License: Apache-2.0
 Project-URL: Homepage, https://github.com/cai4cai/defacing_pipeline
 Project-URL: Repository, https://github.com/cai4cai/defacing_pipeline
 Keywords: MRI,defacing,anonymisation,skull-stripping,neuroimaging,NER,text-anonymization
@@ -13,7 +13,6 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-License-File: LICENSE.md
 Requires-Dist: nibabel>=4.0
 Requires-Dist: numpy<2,>=1.22
 Requires-Dist: scipy>=1.9
@@ -27,7 +26,8 @@ Requires-Dist: faker>=18.0
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
-Dynamic: license-file
+Provides-Extra: ct
+Requires-Dist: TotalSegmentator; extra == "ct"
 
 # caideface
 
@@ -330,4 +330,4 @@ If you use the text anonymisation (NER + HIPS), please also cite:
 
 ## License
 
-This project is licensed under the MIT License -- see the [LICENSE](LICENSE.md) file for details.
+This project is licensed under the Apache License 2.0 -- see the [LICENSE](https://github.com/cai4cai/defacing_pipeline/blob/main/LICENSE) file for details.

{caideface-0.3.0 → caideface-0.3.2}/src/caideface.egg-info/SOURCES.txt

@@ -1,4 +1,3 @@
-LICENSE.md
 README.md
 pyproject.toml
 src/caideface/__init__.py
@@ -15,6 +14,7 @@ src/caideface.egg-info/dependency_links.txt
 src/caideface.egg-info/entry_points.txt
 src/caideface.egg-info/requires.txt
 src/caideface.egg-info/top_level.txt
+src/caideface/data/ct_face_mask.nii.gz
 src/caideface/data/mni_icbm152_t1_tal_nlin_sym_55_ext_brain_only.nii.gz
 src/caideface/data/t1_mask.nii.gz
 src/caideface/data/ner_model/config.cfg

{caideface-0.3.0 → caideface-0.3.2}/src/caideface.egg-info/requires.txt

@@ -9,6 +9,9 @@ hd-bet
 spacy>=3.5
 faker>=18.0
 
+[ct]
+TotalSegmentator
+
 [dev]
 pytest
 ruff

caideface-0.3.0/LICENSE.md

@@ -1,21 +0,0 @@
-# MIT License
-
-Copyright (c) 2025-present cai4cai
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.