crop_dwi_to_mask_bbox.py

#!/usr/bin/env python3
"""
Crop a 4D DWI NIfTI to the bounding box of a (non-zero) mask, with optional padding.

This script is extracted from preprocessing.sh to keep the bash pipeline clean.

Outputs (defaults match the original inline script):
  - mask_crop.nii.gz
  - dwmri_crop.nii.gz
  - crop_bbox.json
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import nibabel as nib
import numpy as np


def crop_to_mask_bbox(
    *,
    dwi_path: Path,
    mask_path: Path,
    out_mask: Path,
    out_dwi: Path,
    out_json: Path,
    pad: int = 1,
) -> None:
    if pad < 0:
        raise ValueError("pad must be >= 0")
    if not dwi_path.exists():
        raise FileNotFoundError(f"Missing expected file: {dwi_path}")
    if not mask_path.exists():
        raise FileNotFoundError(f"Missing expected file: {mask_path}")

    m = nib.load(str(mask_path))
    md = m.get_fdata() > 0
    coords = np.array(np.where(md))
    if coords.size == 0:
        raise RuntimeError("Mask is empty - cannot crop.")

    mins = coords.min(axis=1)
    maxs = coords.max(axis=1)

    mins = np.maximum(mins - pad, 0)
    maxs = np.minimum(maxs + pad, np.array(md.shape) - 1)

    slices = tuple(slice(int(mins[i]), int(maxs[i]) + 1) for i in range(3))

    # Crop mask (adjust affine translation so world coords remain correct)
    m_crop = md[slices].astype(np.uint8)
    A = m.affine.copy()
    A[:3, 3] = A[:3, 3] + A[:3, :3] @ mins.astype(float)
    nib.save(nib.Nifti1Image(m_crop, A), str(out_mask))

    # Crop DWI (4D)
    d = nib.load(str(dwi_path))
    d_data = d.get_fdata(dtype=np.float32)[slices + (slice(None),)]
    Ad = d.affine.copy()
    Ad[:3, 3] = Ad[:3, 3] + Ad[:3, :3] @ mins.astype(float)
    nib.save(nib.Nifti1Image(d_data, Ad, d.header), str(out_dwi))

    bbox = {"mins": mins.tolist(), "maxs": maxs.tolist(), "pad": int(pad)}
    out_json.write_text(json.dumps(bbox, indent=2))
    print("Cropped bbox:", bbox)


def main() -> None:
    p = argparse.ArgumentParser(
        description="Crop DWI to the mask bounding box (with padding)."
    )
    p.add_argument("--dwi", required=True, type=Path, help="Input 4D DWI NIfTI (e.g. dwi_full.nii.gz)")
    p.add_argument("--mask", required=True, type=Path, help="Input mask NIfTI (non-zero defines crop box)")
    p.add_argument("--pad", type=int, default=1, help="Padding in voxels to add to bbox (default: 1)")
    p.add_argument("--out-mask", type=Path, default=Path("mask_crop.nii.gz"))
    p.add_argument("--out-dwi", type=Path, default=Path("dwmri_crop.nii.gz"))
    p.add_argument("--out-json", type=Path, default=Path("crop_bbox.json"))
    args = p.parse_args()

    crop_to_mask_bbox(
        dwi_path=args.dwi,
        mask_path=args.mask,
        out_mask=args.out_mask,
        out_dwi=args.out_dwi,
        out_json=args.out_json,
        pad=args.pad,
    )


if __name__ == "__main__":
    main()