atlas_preprocess.sh

#!/bin/bash

# Exit immediately if a command exits with a non-zero status
set -e

# Default values
T1W_IMG=""
LESION_MASK=""
OUT_DIR=""
USE_GPU=0
CUDA_DEV=""
VERBOSE=0

# Display usage
usage() {
    echo "Usage: $0 -i <T1w_image> -o <output_directory> [options]"
    echo ""
    echo "Options:"
    echo "  -i, --img       Path to T1w image (required)"
    echo "  -o, --outdir    Path to output directory (required)"
    echo "  -m, --mask      Path to lesion mask (optional)"
    echo "  -g, --gpu       Flag to specify if GPU should be used"
    echo "  -c, --cuda      Flag to indicate which GPU to use (e.g. 0)"
    echo "  -v, --verbose   Keep all files (1). Default is 0 (keeps only core files)."
    echo "  -h, --help      Display this help message"
    exit 1
}

# Parse command line arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -i|--img) T1W_IMG="$2"; shift ;;
        -o|--outdir) OUT_DIR="$2"; shift ;;
        -m|--mask) LESION_MASK="$2"; shift ;;
        -g|--gpu) USE_GPU=1 ;;
        -c|--cuda) CUDA_DEV="$2"; shift ;;
        -v|--verbose) 
            if [[ "$2" == "0" || "$2" == "1" ]]; then
                VERBOSE="$2"
                shift
            else
                VERBOSE=1
            fi
            ;;
        -h|--help) usage ;;
        *) echo "Unknown parameter passed: $1"; usage ;;
    esac
    shift
done

if [ -z "$T1W_IMG" ] || [ -z "$OUT_DIR" ]; then
    echo "Error: T1w image (-i) and output directory (-o) are required."
    usage
fi

# Load configuration
CONFIG_FILE="$(dirname "$0")/config.sh"
if [ ! -f "$CONFIG_FILE" ]; then
    echo "Error: Configuration file config.sh not found in the script directory!"
    exit 1
fi
source "$CONFIG_FILE"

# Verify required executables exist
if [ ! -x "${ANTS_PATH}/N4BiasFieldCorrection" ] || [ ! -x "${ANTS_PATH}/ImageMath" ]; then
    echo "Error: ANTs binaries not found in ${ANTS_PATH}. Please check config.sh."
    exit 1
fi

if [ -z "$FREESURFER_HOME" ] || [ ! -x "${FREESURFER_HOME}/bin/mri_synthmorph" ]; then
    echo "Error: FreeSurfer binaries not found under ${FREESURFER_HOME}. Please check FREESURFER_HOME in config.sh."
    exit 1
fi

if [ ! -f "$MNI_TEMPLATE" ]; then
    echo "Error: MNI template not found at '${MNI_TEMPLATE}'. Please check config.sh."
    exit 1
fi

# Set up FreeSurfer environment temporarily for this script run
export FREESURFER_HOME="$FREESURFER_HOME"
source "${FREESURFER_HOME}/SetUpFreeSurfer.sh" >/dev/null 2>&1 || true

# Make sure OUT_DIR exists
mkdir -p "$OUT_DIR"

# Set up GPU flags
GPU_FLAG=""
if [ "$USE_GPU" -eq 1 ]; then
    GPU_FLAG="-g"
    if [ -n "$CUDA_DEV" ]; then
        export CUDA_VISIBLE_DEVICES="$CUDA_DEV"
    fi
fi

# Extract base filename without extension
base_name=$(basename "$T1W_IMG")
base_name="${base_name%%.*}"
base_name="${base_name%_T1w}"

# Extract subject and session from filename for BIDS format
SUB_ID=$(echo "$base_name" | grep -oE 'sub-[a-zA-Z0-9]+' || echo "sub-unknown")
SES_ID=$(echo "$base_name" | grep -oE 'ses-[a-zA-Z0-9]+' || echo "ses-unknown")

# Define step-specific directories
SYNTHSTRIP_DIR="${OUT_DIR}/derivatives/synthstrip/${SUB_ID}/${SES_ID}"
N4_DIR="${OUT_DIR}/derivatives/ants_N4BiasFieldCorrection/${SUB_ID}/${SES_ID}"
SYNTHMORPH_DIR="${OUT_DIR}/derivatives/synthmorph_affine/${SUB_ID}/${SES_ID}"

# Make sure all directories exist
mkdir -p "$SYNTHSTRIP_DIR" "$N4_DIR" "$SYNTHMORPH_DIR"

# Define all intermediate and final output paths
AFFINE_XFM="${SYNTHMORPH_DIR}/${base_name}_orig_to_MNI152NLin2009aSym.lta"
BRAIN_MASK="${SYNTHSTRIP_DIR}/${base_name}_proc-synthstrip_space-orig_desc-T1brain_mask.nii.gz"
MNI_MASK="${SYNTHSTRIP_DIR}/${base_name}_proc-synthstrip_space-MNI152NLin2009aSym_desc-T1brain_mask.nii.gz"

N4_IMG="${N4_DIR}/${base_name}_desc-N4_T1w.nii.gz"
N4_MNI_IMG="${SYNTHMORPH_DIR}/${base_name}_space-MNI152NLin2009aSym_desc-N4_T1w.nii.gz"
N4_MASKED="${SYNTHMORPH_DIR}/${base_name}_space-MNI152NLin2009aSym_desc-N4masked_T1w.nii.gz"
DYNAMIC_RANGE="${SYNTHMORPH_DIR}/${base_name}_space-MNI152NLin2009aSym_desc-dynamicrange_T1w.nii.gz"
BRAIN_NORM="${SYNTHMORPH_DIR}/${base_name}_space-MNI152NLin2009aSym_desc-brainnorm_T1w.nii.gz"

echo "------------------------------------------------------"
echo "Starting atlas_preprocess pipeline for: $T1W_IMG"
echo "Output directory: $OUT_DIR"
if [ "$USE_GPU" -eq 1 ]; then echo "GPU: Enabled (CUDA_VISIBLE_DEVICES=${CUDA_DEV})"; else echo "GPU: Disabled"; fi
echo "Verbose mode: $VERBOSE"
echo "------------------------------------------------------"

# Step 1: Synthmorph affine registration to MNI space
echo "[1/6] Running mri_synthmorph affine registration to MNI space..."
mri_synthmorph register $GPU_FLAG -a "$AFFINE_XFM" "$T1W_IMG" "$MNI_TEMPLATE"

# Step 2: Synthstrip to get brain mask
echo "[2/6] Running mri_synthstrip to extract brain mask..."
mri_synthstrip -i "$T1W_IMG" -m "$BRAIN_MASK" $GPU_FLAG

# Apply affine to nearest-neighbor align the brain mask to MNI
echo "      Applying affine transform to brain mask to align to MNI..."
mri_synthmorph apply -m nearest "$AFFINE_XFM" "$BRAIN_MASK" "$MNI_MASK"

# Step 3: N4 Bias field correction
echo "[3/6] Running N4 Bias Field Correction..."
"${ANTS_PATH}/N4BiasFieldCorrection" -i "$T1W_IMG" -o "$N4_IMG" -x "$BRAIN_MASK"

# Step 4: Apply synthmorph affine to N4 corrected image
echo "[4/6] Applying affine transform to N4-corrected T1w image..."
mri_synthmorph apply "$AFFINE_XFM" "$N4_IMG" "$N4_MNI_IMG"

if [ -n "$LESION_MASK" ]; then
    if [ -f "$LESION_MASK" ]; then
        MNI_LESION_MASK="${SYNTHMORPH_DIR}/${base_name}_space-MNI152NLin2009aSym_desc-lesion_mask.nii.gz"
        echo "      Lesion mask provided. Applying affine transform to align lesion mask to MNI..."
        mri_synthmorph apply -m nearest "$AFFINE_XFM" "$LESION_MASK" "$MNI_LESION_MASK"
    else
        echo "Warning: Provided lesion mask '$LESION_MASK' does not exist."
    fi
fi

# Step 5: Mask the MNI image
echo "[5/6] Masking N4-corrected MNI image with MNI brain mask..."
"${ANTS_PATH}/ImageMath" 3 "$N4_MASKED" m "$N4_MNI_IMG" "$MNI_MASK"

# Step 6: Dynamic range truncate and rescale intensity
echo "[6/6] Truncating dynamic range and rescaling intensity (0-100)..."
"${ANTS_PATH}/ImageMath" 3 "$DYNAMIC_RANGE" TruncateImageIntensity "$N4_MASKED" 0.005 0.995
"${ANTS_PATH}/ImageMath" 3 "$BRAIN_NORM" RescaleImage "$DYNAMIC_RANGE" 0 100

# Cleanup intermediate files if not verbose
if [ "$VERBOSE" -eq 0 ]; then
    echo "------------------------------------------------------"
    echo "Cleaning up intermediate files (verbose=0)..."
    rm -f "$N4_MNI_IMG"
    rm -f "$N4_MASKED"
    rm -f "$DYNAMIC_RANGE"
fi

echo "------------------------------------------------------"
echo "Pipeline completed successfully!"
echo "Final normalized image: $BRAIN_NORM"