BML_Burn_Scriptv2.py

#!/usr/bin/env python3
"""
Burning Script for Brain Modulation Lab
Author: Clemens Neudorfer
Date: October 18, 2024

This script identifies nii files converted into native space based on the prefix w*, aggregates them and burns them into a 
native space T1 scan. Native nii files can be represented as binary or continuous.

Usage:
    python3 burn_karawun.py [--continuous true|false]

Arguments:
    --continuous    Optional flag to enable continuous intensity modulation.
                    If not provided, binary modulation is applied.

Prerequisites:
    - Ensure that the Conda environment 'KarawunEnv' is activated before running.
    - FSL tools (e.g., fslmaths, fslstats) must be installed and accessible.
    - importTractography should be available in the system PATH.
"""

import os
import glob
import shutil
import subprocess
import argparse
import sys

def get_intensity_range(filepath):
    """
    Retrieve the minimum and maximum intensity values of a NIfTI image.

    Args:
        filepath (str): Path to the NIfTI file.

    Returns:
        tuple: (min_intensity, max_intensity) as floats.

    Raises:
        subprocess.CalledProcessError: If the fslstats command fails.
        ValueError: If the output cannot be parsed into two floats.
    """
    cmd = f"fslstats {filepath} -R"
    try:
        output = subprocess.check_output(cmd, shell=True, text=True).strip()
        min_intensity, max_intensity = map(float, output.split())
        return min_intensity, max_intensity
    except subprocess.CalledProcessError as e:
        print(f"Error executing command: {cmd}\n{e}")
        sys.exit(1)
    except ValueError:
        print(f"Unexpected output from fslstats: '{output}'")
        sys.exit(1)

def create_directory(path):
    """
    Create a directory if it does not already exist.

    Args:
        path (str): Path to the directory.
    """
    os.makedirs(path, exist_ok=True)

def process_atlases(atlas_files, seedmasks_dir, burn_nuclei_dir, continuous, t1_min, t1_max):
    """
    Process atlas files by thresholding, masking, and merging with the T1 image.

    Args:
        atlas_files (list): List of atlas file paths.
        seedmasks_dir (str): Directory to store seed masks.
        burn_nuclei_dir (str): Directory to store burned nuclei.
        continuous (bool): Flag indicating whether to apply continuous modulation.
        t1_min (float): Minimum intensity of the T1 image.
        t1_max (float): Maximum intensity of the T1 image.

    Returns:
        str: Path to the final burned T1 image.
    """
    # Define intensity levels for binary modulation
    intensities = [
        t1_min + 1,
        t1_max,
        (t1_max + t1_min) * 3 / 5,
        (t1_max + t1_min) * 2 / 5,
        (t1_max + t1_min) * 3 / 5,
        (t1_max + t1_min) * 4 / 5
    ]

    # Initialize the burned T1 image
    anat_t1_burned = os.path.join(burn_nuclei_dir, 'anat_t1_burned.nii.gz')
    shutil.copyfile(os.path.join(os.getcwd(), 'anat_t1.nii'), anat_t1_burned)

    for idx, atlas in enumerate(atlas_files):
        atlas_basename = os.path.basename(atlas)
        atlas_path = os.path.join(seedmasks_dir, atlas_basename)

        # Threshold atlas at 0.5
        if continuous:
            thresh_cmd = f"fslmaths {atlas_path} -thr 0.5 {atlas_path}"
        else:
            thresh_cmd = f"fslmaths {atlas_path} -thr 0.5 -bin {atlas_path}"

        subprocess.run(thresh_cmd, shell=True, check=True)

        # Create inverted mask
        mask_inverted = os.path.join(seedmasks_dir, f"mask_{idx}.nii.gz")
        invert_cmd = f"fslmaths {atlas_path} -bin -sub 1 -mul -1 {mask_inverted}"
        subprocess.run(invert_cmd, shell=True, check=True)

        # Create binary mask
        mask_binary = os.path.join(seedmasks_dir, f"mask2_{idx}.nii.gz")
        mask_cmd = f"fslmaths {atlas_path} -bin {mask_binary}"
        subprocess.run(mask_cmd, shell=True, check=True)

        # Apply mask to the burned T1 image
        masked_t1_cmd = f"fslmaths {anat_t1_burned} -mul {mask_inverted} {anat_t1_burned}"
        subprocess.run(masked_t1_cmd, shell=True, check=True)

        if continuous:
            # Get intensity range of the atlas
            atlas_min, atlas_max = get_intensity_range(atlas_path)

            # Calculate scale factor and offset
            scale_factor = abs((t1_max - t1_min) / (atlas_max - atlas_min))
            offset = abs(t1_min - (atlas_min * scale_factor))

            # Apply scaling and offset to create a gradient
            set_gradient_cmd = f"fslmaths {atlas_path} -mul {scale_factor} -add {offset} {atlas_path}"
            subprocess.run(set_gradient_cmd, shell=True, check=True)

            # Apply binary mask to the scaled atlas
            mask_atlas_cmd = f"fslmaths {atlas_path} -mul {mask_binary} {atlas_path}"
            subprocess.run(mask_atlas_cmd, shell=True, check=True)
        else:
            try:
                intensity = intensities[idx]
            except IndexError:
                print(f"Insufficient intensity levels for atlas index {idx}.")
                sys.exit(1)
            # Apply binary mask with predefined intensity
            set_intensity_cmd = f"fslmaths {atlas_path} -bin -mul {intensity} {atlas_path}"
            subprocess.run(set_intensity_cmd, shell=True, check=True)

        # Merge the current atlas into the burned T1 image
        merge_cmd = f"fslmaths {anat_t1_burned} -add {atlas_path} {anat_t1_burned}"
        subprocess.run(merge_cmd, shell=True, check=True)

    return anat_t1_burned

def burn_combined_cluster(subjects_dir, burn_nuclei_dir):
    """
    Burn the combined cluster into a single T1 image using importTractography.

    Args:
        subjects_dir (str): Path to the subjects directory.
        burn_nuclei_dir (str): Directory containing the burned nuclei.
    """
    dicom_input = os.path.join(subjects_dir, 'dummy00001.dcm')
    dicom_output_dir = os.path.join(burn_nuclei_dir, 'DICOM')
    anat_burned = os.path.join(burn_nuclei_dir, 'anat_t1_burned.nii.gz')

    # Ensure DICOM output directory exists
    create_directory(dicom_output_dir)

    burn_cmd = (
        f"importTractography -d {dicom_input} "
        f"-o {dicom_output_dir}/ "
        f"-n {anat_burned}"
    )
    try:
        subprocess.run(burn_cmd, shell=True, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error executing importTractography: {e}")
        sys.exit(1)

def main():
    """Main function to execute the burning script."""
    # Set up argument parser
    parser = argparse.ArgumentParser(
        description='Burn atlases into anatomical T1 images with optional continuous modulation.'
    )
    parser.add_argument(
        '--continuous',
        type=str,
        default='false',
        help='Enable continuous intensity modulation (true|false). Default is false.'
    )

    # Parse arguments
    args = parser.parse_args()
    continuous = args.continuous.lower() == 'true'

    # Define working directories
    current_dir = os.getcwd()
    seedmasks_dir = os.path.join(current_dir, 'seedmasks')
    burn_nuclei_dir = os.path.join(current_dir, 'burn_nuclei')

    print("Initiating burning procedure.")

    # Create necessary directories
    create_directory(seedmasks_dir)
    create_directory(burn_nuclei_dir)

    # Retrieve atlas files (assuming they start with 'w')
    atlas_pattern = os.path.join(current_dir, 'w*')
    atlas_files = sorted(glob.glob(atlas_pattern))

    if not atlas_files:
        print(f"No atlas files found with pattern: {atlas_pattern}")
        sys.exit(1)

    # Copy atlas files to seedmasks directory
    copied_atlases = []
    for atlas in atlas_files:
        destination = os.path.join(seedmasks_dir, os.path.basename(atlas))
        try:
            shutil.copyfile(atlas, destination)
            copied_atlases.append(destination)
        except IOError as e:
            print(f"Failed to copy {atlas} to {destination}: {e}")
            sys.exit(1)

    # Apply initial threshold to copied atlases
    for atlas in copied_atlases:
        thresh_cmd = f"fslmaths {atlas} -thr 0.3 {atlas}"
        try:
            subprocess.run(thresh_cmd, shell=True, check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error thresholding atlas {atlas}: {e}")
            sys.exit(1)

    # Get intensity range of the anatomical T1 image
    anat_t1_path = os.path.join(current_dir, 'anat_t1.nii')
    if not os.path.isfile(anat_t1_path):
        print(f"Anatomical T1 image not found at: {anat_t1_path}")
        sys.exit(1)

    t1_min, t1_max = get_intensity_range(anat_t1_path)

    # Process atlases and merge into the burned T1 image
    anat_burned = process_atlases(
        atlas_files=copied_atlases,
        seedmasks_dir=seedmasks_dir,
        burn_nuclei_dir=burn_nuclei_dir,
        continuous=continuous,
        t1_min=t1_min,
        t1_max=t1_max
    )

    # Burn the combined cluster into the final T1 image
    burn_combined_cluster(current_dir, burn_nuclei_dir)

    # Delete seedmasks folder after completion
    shutil.rmtree(seedmasks_dir)

    print("Burning process completed successfully.")

if __name__ == '__main__':
    main()