preprocessing.sh

#!/usr/bin/env bash
# make generalizability figure
# for each dataset - grab HC, grab non-HC from prequal
# # for each:
# 0) Convert DWI to .mif once (so MRtrix tools are happy)
# 1) Mean b0 (for BET)
# 2) Two masks + combine
# 3) Crop DWI to mask bbox (+1 voxel)
# 4) Upsample cropped DWI and mask to 1mm (your preference)
# 5) Convert cropped+1mm DWI to .mif with correct strides
# 6) Compute response + WM FOD (single-shell vs multi-shell)
# 7) Sanity check: make sure fODF has 45 coefficients
# # copy to hendrix
# # Running inference on Hendrix (v3 fit+pad176)
# # QA inference
# 8) optional tractography

# inputs have to nii.gz, bval, bvec (preprocessed data) - any resolution, any b-values
# options: default is to resample to 1mm
# options: tractography is run if enabled
# QA: bundles/parcellation (tractography if run)
# final option (low low priority): putting back in regular space

set -euo pipefail

usage() {
  cat <<'EOF'
Usage:
  singularity run -B /path/to/input:/INPUT -B /path/to/output:/OUTPUT container.sif [OPTIONS]

The container reads from /INPUT and writes to /OUTPUT.
Bind your host directories to these paths.

Stages (composable):
  --preprocess      Run preprocessing only (up to creating wmfod.nii.gz).
  --inference       Run CPU inference (BundleParc). Requires wmfod.nii.gz.
  --QA              Run QA rendering (PNGs + PDF report) for bundle segmentations.
  --tracking        Run tractography (.tck) from predicted label maps.
  --tract_qa        Run QA rendering for tractography (requires Tractography/).

Options:
  --nb_pts N        Number of label divisions per bundle (default 10).
  --keep_biggest_blob  Keep only the largest connected component per bundle
                       (overrides default min_blob_size=0).
  -h, --help        Show this help and exit.

Assumptions:
  - Processing ONE subject at a time.
  - /INPUT must contain exactly ONE .nii.gz, ONE .bval, and ONE .bvec file.

Outputs (written to /OUTPUT):
  - Intermediate + final fODF: wmfod.nii.gz
  - Segmentations: Brainstem-Segmentation/*.nii.gz
  - QA PNGs: QA/*.png, QA_tract/*.png
  - QA PDFs: PDF/qa_report.pdf, PDF/tract_qa_report.pdf

Notes:
  - If no stage flags are given, --preprocess and --inference run by default.
EOF
}

need_cmd() { command -v "$1" >/dev/null 2>&1 || { echo "ERROR: missing required command: $1" >&2; exit 127; }; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Singularity runs containers with a read-only root filesystem.
# Force MRtrix (and other tools) to use a writable temp directory.
export TMPDIR="${TMPDIR:-/tmp}"
export MRTRIX_TMPFILE_DIR="${MRTRIX_TMPFILE_DIR:-$TMPDIR}"
mkdir -p "$MRTRIX_TMPFILE_DIR" 2>/dev/null || true

safe_deactivate_venv() {
  # Only virtualenv-style activation defines a deactivate() shell function.
  if [[ -n "${VIRTUAL_ENV:-}" ]] && [[ "$(type -t deactivate || true)" == "function" ]]; then
    deactivate || true
  fi
}

safe_source_venv() {
  local activate_path="$1"
  [[ -f "$activate_path" ]] || return 1
  safe_deactivate_venv
  # shellcheck disable=SC1090
  source "$activate_path"
}

activate_scil_env_if_present() {
  # Needed for nibabel in crop script (and any other Python helpers).
  if [[ -n "${VIRTUAL_ENV:-}" ]] && [[ -f "${VIRTUAL_ENV}/bin/activate" ]]; then
    safe_source_venv "${VIRTUAL_ENV}/bin/activate" || true
  elif [[ -f "/opt/venvs/scil_env/.venv/bin/activate" ]]; then
    safe_source_venv "/opt/venvs/scil_env/.venv/bin/activate" || true
  fi
}

source_fsl_if_present() {
  export FSLDIR="${FSLDIR:-/root/fsl}"
  if [[ -f "${FSLDIR}/etc/fslconf/fsl.sh" ]]; then
    # shellcheck disable=SC1090
    source "${FSLDIR}/etc/fslconf/fsl.sh"
  fi
  export FSLOUTPUTTYPE="${FSLOUTPUTTYPE:-NIFTI_GZ}"
}

parse_args() {
  DO_PREPROCESS=0
  DO_INFERENCE=0
  DO_QA=0
  DO_TRACKING=0
  DO_TRACT_QA=0
  NB_PTS=10
  KEEP_BIGGEST_BLOB=0

  while [[ $# -gt 0 ]]; do
    case "$1" in
      # Composable stage flags
      --preprocess)    DO_PREPROCESS=1; shift ;;
      --inference)     DO_INFERENCE=1; shift ;;
      --QA|--qa)       DO_QA=1; shift ;;
      --tracking)      DO_TRACKING=1; shift ;;
      --tract_qa|--tract-qa|--tract_QA)
                       DO_TRACT_QA=1; shift ;;

      # Backward-compatible aliases
      --preprocessing) DO_PREPROCESS=1; shift ;;

      # Options
      --nb_pts|--nb-pts)
        NB_PTS="$2"; shift 2 ;;
      --keep_biggest_blob|--keep-biggest-blob)
        KEEP_BIGGEST_BLOB=1; shift ;;

      -h|--help) usage; exit 0 ;;
      -*)
        echo "ERROR: unknown option: $1" >&2
        usage
        exit 2
        ;;
      *)
        echo "ERROR: unexpected positional argument: $1" >&2
        echo "  Input/output are fixed: /INPUT and /OUTPUT. Bind your host dirs there." >&2
        usage
        exit 2
        ;;
    esac
  done

  # Default behavior if no stage flags were provided:
  if [[ "${DO_PREPROCESS}" == "0" && "${DO_INFERENCE}" == "0" && "${DO_QA}" == "0" && "${DO_TRACKING}" == "0" && "${DO_TRACT_QA}" == "0" ]]; then
    DO_PREPROCESS=1
    DO_INFERENCE=1
    DO_QA=1
  fi
}

# Fixed container paths — bind your host directories here.
SUBJ_DIR="/INPUT"
PROC="/OUTPUT"

resolve_inputs() {
  if [[ ! -d "$SUBJ_DIR" ]]; then
    echo "ERROR: /INPUT is not mounted. Bind your input directory with: -B /host/path:/INPUT" >&2
    exit 2
  fi
  if [[ ! -d "$PROC" ]]; then
    echo "ERROR: /OUTPUT is not mounted. Bind your output directory with: -B /host/path:/OUTPUT" >&2
    exit 2
  fi
}

run_preprocessing() {
  activate_scil_env_if_present
  source_fsl_if_present

  # ── Auto-detect input files by suffix ──
  shopt -s nullglob
  local -a nifti_files=("$SUBJ_DIR"/*.nii.gz)
  local -a bval_files=("$SUBJ_DIR"/*.bval)
  local -a bvec_files=("$SUBJ_DIR"/*.bvec)
  shopt -u nullglob

  if [[ ${#nifti_files[@]} -eq 0 ]]; then
    echo "ERROR: no .nii.gz file found in $SUBJ_DIR" >&2; exit 2
  elif [[ ${#nifti_files[@]} -gt 1 ]]; then
    echo "ERROR: multiple .nii.gz files found in $SUBJ_DIR (expected exactly 1):" >&2
    printf '  %s\n' "${nifti_files[@]}" >&2; exit 2
  fi
  if [[ ${#bval_files[@]} -eq 0 ]]; then
    echo "ERROR: no .bval file found in $SUBJ_DIR" >&2; exit 2
  elif [[ ${#bval_files[@]} -gt 1 ]]; then
    echo "ERROR: multiple .bval files found in $SUBJ_DIR (expected exactly 1):" >&2
    printf '  %s\n' "${bval_files[@]}" >&2; exit 2
  fi
  if [[ ${#bvec_files[@]} -eq 0 ]]; then
    echo "ERROR: no .bvec file found in $SUBJ_DIR" >&2; exit 2
  elif [[ ${#bvec_files[@]} -gt 1 ]]; then
    echo "ERROR: multiple .bvec files found in $SUBJ_DIR (expected exactly 1):" >&2
    printf '  %s\n' "${bvec_files[@]}" >&2; exit 2
  fi

  DWI_NII="${nifti_files[0]}"
  BVAL="${bval_files[0]}"
  BVEC="${bvec_files[0]}"
  echo "Detected inputs:"
  echo "  NIfTI: $DWI_NII"
  echo "  bval:  $BVAL"
  echo "  bvec:  $BVEC"

  need_cmd mrconvert
  need_cmd dwiextract
  need_cmd mrmath
  need_cmd dwi2mask
  need_cmd mrcalc
  need_cmd mrgrid
  need_cmd mrinfo
  need_cmd bet
  need_cmd python3

  echo "===================="
  echo "Stage: preprocess"
  echo "Input:  $SUBJ_DIR"
  echo "Output: $PROC"
  echo "===================="

  ##### 0) Convert DWI to .mif once (so MRtrix tools are happy)
  mrconvert "$DWI_NII" "$PROC/dwi_full.mif" -fslgrad "$BVEC" "$BVAL" -strides 1,2,3,4 -force
  mrconvert "$PROC/dwi_full.mif" "$PROC/dwi_full.nii.gz" -force

  ##### 1) Mean b0 (for BET)
  dwiextract "$PROC/dwi_full.mif" -bzero - | mrmath - mean "$PROC/meanb0.mif" -axis 3 -force
  mrconvert "$PROC/meanb0.mif" "$PROC/meanb0.nii.gz" -force

  ##### 2) Two masks + combine
  dwi2mask "$PROC/dwi_full.mif" "$PROC/mask_dwi2mask.mif" -force
  mrconvert "$PROC/mask_dwi2mask.mif" "$PROC/mask_dwi2mask.nii.gz" -force
  bet "$PROC/meanb0.nii.gz" "$PROC/meanb0_bet" -m -R -f 0.25 -n
  mrconvert "$PROC/meanb0_bet_mask.nii.gz" "$PROC/mask_bet.mif" -datatype bit -force
  mrcalc "$PROC/mask_dwi2mask.mif" "$PROC/mask_bet.mif" -max "$PROC/mask_union.mif" -force
  mrconvert "$PROC/mask_union.mif" "$PROC/mask_union.nii.gz" -force

  ##### 3) Crop DWI to mask bbox (+1 voxel)
  (
    cd "$PROC"
    python3 "$SCRIPT_DIR/crop_dwi_to_mask_bbox.py" \
      --dwi "dwi_full.nii.gz" \
      --mask "mask_union.nii.gz" \
      --pad 1 \
      --out-mask "mask_crop.nii.gz" \
      --out-dwi "dwmri_crop.nii.gz" \
      --out-json "crop_bbox.json"
  )

  ##### 4) Upsample cropped DWI and mask to 1mm
  mrgrid "$PROC/dwmri_crop.nii.gz" regrid "$PROC/dwmri_crop_1mm.nii.gz" -voxel 1 -force
  mrgrid "$PROC/mask_crop.nii.gz"  regrid "$PROC/mask_crop_1mm.nii.gz"  -voxel 1 -interp nearest -force

  ##### 5) Convert cropped+1mm DWI to .mif with correct strides
  mrconvert "$PROC/dwmri_crop_1mm.nii.gz" "$PROC/dwi.mif" -fslgrad "$BVEC" "$BVAL" -strides 1,2,3,4 -force
  mrconvert "$PROC/mask_crop_1mm.nii.gz" "$PROC/mask.nii.gz" -datatype bit -force

  ###### 6) Compute response + WM FOD (single-shell vs multi-shell)
  n=$(mrinfo "$PROC/dwi.mif" -shell_bvalues | tr -d , | awk -v thr=50 '{for(i=1;i<=NF;i++) if($i>thr) c++} END{print c+0}')
  if [[ "$n" -gt 1 ]]; then
    echo "Shells detected: multishell"
    dwi2response dhollander "$PROC/dwi.mif" \
      "$PROC/response_wm.txt" "$PROC/response_gm.txt" "$PROC/response_csf.txt" \
      -mask "$PROC/mask.nii.gz" -force
    dwi2fod msmt_csd "$PROC/dwi.mif" \
      "$PROC/response_wm.txt" "$PROC/wmfod.nii.gz" \
      "$PROC/response_gm.txt" "$PROC/gm.mif" \
      "$PROC/response_csf.txt" "$PROC/csf.mif" \
      -mask "$PROC/mask.nii.gz" -force
  else
    echo "Shells detected: singleshell"
    dwi2response tournier "$PROC/dwi.mif" "$PROC/response_wm.txt" -mask "$PROC/mask.nii.gz" -force
    dwi2fod csd "$PROC/dwi.mif" "$PROC/response_wm.txt" "$PROC/wmfod.nii.gz" -mask "$PROC/mask.nii.gz" -force
  fi

  echo "Done. Output: $PROC/wmfod.nii.gz"
}

run_inference_cpu() {
  echo "===================="
  echo "Stage: inference (CPU)"
  echo "Input:  $SUBJ_DIR"
  echo "Output: $PROC"
  echo "===================="

  need_cmd python3
  export CUDA_VISIBLE_DEVICES=""
  export COMET_MODE="${COMET_MODE:-DISABLED}"

  CKPT_PATH="/opt/pipeline/checkpoints/model.ckpt"
  IN_FODF="$PROC/wmfod.nii.gz"
  OUT_DIR="$PROC/Brainstem-Segmentation"
  mkdir -p "$OUT_DIR"

  if [[ ! -f "$IN_FODF" ]]; then
    echo "ERROR: missing expected fODF: $IN_FODF (run with --preprocess first, or include --preprocess)" >&2
    exit 2
  fi

  # Run inference inside BundleParc venv (preferred)
  BUNDLEPARC_PY="/opt/bundleparc/.venv/bin/python"
  if [[ -x "$BUNDLEPARC_PY" ]]; then
    # Activate the venv for env vars, but always execute the venv python explicitly.
    if [[ -f "/opt/bundleparc/.venv/bin/activate" ]]; then
      safe_source_venv "/opt/bundleparc/.venv/bin/activate"
    fi

    if ! "$BUNDLEPARC_PY" -c "import scilpy" >/dev/null 2>&1; then
      echo "ERROR: BundleParc venv exists but 'scilpy' is not importable." >&2
      echo "  python: $BUNDLEPARC_PY" >&2
      echo "  fix: rebuild the image so /opt/bundleparc/.venv includes dependencies." >&2
      exit 2
    fi

    local blob_args
    if [[ "$KEEP_BIGGEST_BLOB" -eq 1 ]]; then
      blob_args="--keep_biggest_blob"
    else
      blob_args="--min_blob_size 0"
    fi

    "$BUNDLEPARC_PY" /opt/bundleparc/scripts/bundleparc_predict_fitpad.py \
      "$IN_FODF" \
      --checkpoint "$CKPT_PATH" \
      --out_folder "$OUT_DIR" \
      --volume_size 176 \
      $blob_args \
      -f \
      --nb_pts "$NB_PTS"
  else
    echo "WARNING: BundleParc venv python not found at $BUNDLEPARC_PY; falling back to python3" >&2

    local blob_args
    if [[ "$KEEP_BIGGEST_BLOB" -eq 1 ]]; then
      blob_args="--keep_biggest_blob"
    else
      blob_args="--min_blob_size 0"
    fi

    python3 /opt/bundleparc/scripts/bundleparc_predict_fitpad.py \
      "$IN_FODF" \
      --checkpoint "$CKPT_PATH" \
      --out_folder "$OUT_DIR" \
      --volume_size 176 \
      $blob_args \
      -f \
      --nb_pts "$NB_PTS"
  fi

}

run_tracking_only() {
  echo "===================="
  echo "Stage: tracking"
  echo "Output: $PROC"
  echo "===================="

  need_cmd bash
  need_cmd tckgen

  IN_FODF="$PROC/wmfod.nii.gz"
  LABELS_DIR="$PROC/Brainstem-Segmentation"
  TRACK_OUT="$PROC/Tractography"

  if [[ ! -f "$IN_FODF" ]]; then
    echo "ERROR: missing fODF (run --preprocessing or --inference first): $IN_FODF" >&2
    exit 2
  fi
  if [[ ! -d "$LABELS_DIR" ]]; then
    echo "ERROR: missing labels dir (run --inference first): $LABELS_DIR" >&2
    exit 2
  fi

  mkdir -p "$TRACK_OUT"
  bash "$SCRIPT_DIR/tracking_from_labels.sh" \
    --labels_dir "$LABELS_DIR" \
    --fod "$IN_FODF" \
    --out_dir "$TRACK_OUT" \
    --threads auto
}

run_qa_only() {
  echo "===================="
  echo "Stage: QA"
  echo "Output: $PROC"
  echo "===================="

  local SEG_DIR="$PROC/Brainstem-Segmentation"
  local QA_DIR="$PROC/QA"
  local PDF_DIR="$PROC/PDF"
  local MASK_1MM="$PROC/mask_crop_1mm.nii.gz"

  if [[ ! -d "$SEG_DIR" ]]; then
    echo "ERROR: Brainstem-Segmentation folder not found: $SEG_DIR" >&2
    exit 2
  fi
  if [[ ! -f "$MASK_1MM" ]]; then
    echo "ERROR: QA mask not found (expected from preprocessing): $MASK_1MM" >&2
    exit 2
  fi

  mkdir -p "$QA_DIR" "$PDF_DIR"

  local VIZ_PY="/opt/bundle_visualization/.venv/bin/python"
  if [[ ! -x "$VIZ_PY" ]]; then
    echo "WARNING: /opt/bundle_visualization/.venv/bin/python not found; falling back to python3" >&2
    VIZ_PY="python3"
  fi

  # ── Render QA PNGs (batch mode: glass brain computed once) ──
  "$VIZ_PY" /opt/bundle_visualization/bundle_qa.py \
    --mask "$MASK_1MM" \
    --bundle_dir "$SEG_DIR" \
    --out_dir "$QA_DIR" \
    --n_labels "$NB_PTS"

  echo "QA PNGs saved to: $QA_DIR"

  # ── Compile PNGs into PDF report ──
  "$VIZ_PY" /opt/bundle_visualization/compile_qa_pdf.py \
    --qa_dir "$QA_DIR" \
    --out_pdf "$PDF_DIR/qa_report.pdf"

  echo "QA PDF saved to: $PDF_DIR/qa_report.pdf"
}

run_tract_qa_only() {
  echo "===================="
  echo "Stage: Tract QA"
  echo "Output: $PROC"
  echo "===================="

  local TRACT_DIR="$PROC/Tractography"
  local QA_TRACT_DIR="$PROC/QA_tract"
  local PDF_DIR="$PROC/PDF"
  local MASK_1MM="$PROC/mask_crop_1mm.nii.gz"

  if [[ ! -d "$TRACT_DIR" ]]; then
    echo "ERROR: Tractography folder not found: $TRACT_DIR (run --tracking first)" >&2
    exit 2
  fi
  if [[ ! -f "$MASK_1MM" ]]; then
    echo "ERROR: QA mask not found (expected from preprocessing): $MASK_1MM" >&2
    exit 2
  fi

  mkdir -p "$QA_TRACT_DIR" "$PDF_DIR"

  local VIZ_PY="/opt/bundle_visualization/.venv/bin/python"
  if [[ ! -x "$VIZ_PY" ]]; then
    echo "WARNING: /opt/bundle_visualization/.venv/bin/python not found; falling back to python3" >&2
    VIZ_PY="python3"
  fi

  # Render tract QA PNGs
  "$VIZ_PY" /opt/bundle_visualization/tract_qa.py \
    --mask "$MASK_1MM" \
    --tract_dir "$TRACT_DIR" \
    --out_dir "$QA_TRACT_DIR"

  echo "Tract QA PNGs saved to: $QA_TRACT_DIR"

  # Compile into PDF (reuse the same left/right pairing logic)
  "$VIZ_PY" /opt/bundle_visualization/compile_qa_pdf.py \
    --qa_dir "$QA_TRACT_DIR" \
    --out_pdf "$PDF_DIR/tract_qa_report.pdf"

  echo "Tract QA PDF saved to: $PDF_DIR/tract_qa_report.pdf"
}

parse_args "$@"
resolve_inputs

if [[ "${DO_PREPROCESS}" == "1" ]]; then
  run_preprocessing
fi

if [[ "${DO_INFERENCE}" == "1" ]]; then
  run_inference_cpu
fi

if [[ "${DO_TRACKING}" == "1" ]]; then
  run_tracking_only
fi

if [[ "${DO_QA}" == "1" ]]; then
  run_qa_only
fi

if [[ "${DO_TRACT_QA}" == "1" ]]; then
  run_tract_qa_only
fi