From 27a40211666e947e80deff16d98195dc7bbbe7fb Mon Sep 17 00:00:00 2001 From: alessandrofelder Date: Thu, 21 Nov 2024 12:58:44 +0000 Subject: [PATCH 1/3] downsample 7 best molerat images --- examples/molerat/downsample_source_images.py | 101 +++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 examples/molerat/downsample_source_images.py diff --git a/examples/molerat/downsample_source_images.py b/examples/molerat/downsample_source_images.py new file mode 100644 index 0000000..0e58462 --- /dev/null +++ b/examples/molerat/downsample_source_images.py @@ -0,0 +1,101 @@ +import argparse +from pathlib import Path + +import numpy as np +import pandas as pd +from brainglobe_space import AnatomicalSpace +from brainglobe_utils.IO.image import load_any, save_any +from loguru import logger +from skimage import transform + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Downsample source images") + parser.add_argument( + "--template_building_root", + type=str, + help="Path to the template-building root folder. Results will " + "be written to the rawdata subfolder.", + required=True, + ) + parser.add_argument( + "--target_isotropic_resolution", + type=int, + help="Target isotropic resolution", + required=True, + ) + + args = parser.parse_args() + + atlas_forge_molerat_path = Path( + "/media/ceph-neuroinformatics/neuroinformatics/atlas-forge/MoleRat/" + ) + source_data = ( + atlas_forge_molerat_path + / "Mole-rat brain atlas (Fukomys anselli)_MalkemperLab" + ) + source_info_file = atlas_forge_molerat_path / "molerat_brains_info_PM.csv" + + template_building_root = Path(args.template_building_root) + target_isotropic_resolution = int(args.target_isotropic_resolution) + + in_plane_resolution = 20 + out_of_plane_resolution = 20 + in_plane_factor = round(target_isotropic_resolution / in_plane_resolution) + axial_factor = round(target_isotropic_resolution / out_of_plane_resolution) + + template_raw_data = template_building_root / "rawdata" + template_raw_data.mkdir(exist_ok=True, parents=True) + + source_info = pd.read_csv(source_info_file, skiprows=1, header=0) + logger.info(f"Loaded source info with {len(source_info)} entries.") + counter = 0 + for _, sample in source_info.iterrows(): + original_slice_direction = sample["original_slice_direction"] + if original_slice_direction == "horizontal": + source_file = ( + source_data / "HorizontallyImaged" / sample["filename"] + ) + elif original_slice_direction == "sagittal": + source_file = source_data / "SagitallyImaged" / sample["filename"] + else: + raise ValueError( + f"Unexpected slice direction {original_slice_direction}" + ) + assert source_file.exists(), f"File {source_file} not found" + if sample["comments"] != "good": + logger.info(f"Skipping {source_file.name} for now") + continue + + subject_id = sample["subject_id"] + hemisphere = sample["hemisphere"] + subject_folder = ( + template_raw_data / f"sub-{subject_id}_hemi-{hemisphere}" + ) + subject_folder.mkdir(exist_ok=True) + rawdata_filename = ( + f"sub-{subject_id}_" + f"hemi-{hemisphere}_" + f"res-{target_isotropic_resolution}.tif" + ) + + # we can't use our usual transform utils function here, + # because it's not a dask array, + # and we additionally need to mirror and reorient the stack to ASR + assert sample[ + "looks_like_right_hemisphere" + ], f"TODO: flip sample {source_file}" + assert ( + sample["image_orientation"] == "RPI" + ), "Image orientation is not RPI" + stack = load_any(str(source_file)) + downsampled = transform.downscale_local_mean( + stack, (axial_factor, in_plane_factor, in_plane_factor) + ) + downsampled = np.concatenate( + (downsampled, np.flip(downsampled, axis=0)), axis=0 + ) + original_space = AnatomicalSpace("RPI") + downsampled = original_space.map_stack_to("ASR", downsampled) + save_any(downsampled, subject_folder / rawdata_filename) + + logger.info(f"{rawdata_filename} downsampled.") From eaa25e6686895568900acb00717820625aae6823 Mon Sep 17 00:00:00 2001 From: alessandrofelder Date: Fri, 22 Nov 2024 10:16:46 +0000 Subject: [PATCH 2/3] initial test for molerat template --- ...mages.py => 1_downsample_source_images.py} | 32 ++- examples/molerat/2_molerat_prep_lowres.py | 192 ++++++++++++++++++ examples/slurm/build_template_molerat.sh | 30 +++ 3 files changed, 250 insertions(+), 4 deletions(-) rename examples/molerat/{downsample_source_images.py => 1_downsample_source_images.py} (77%) create mode 100644 examples/molerat/2_molerat_prep_lowres.py create mode 100644 examples/slurm/build_template_molerat.sh diff --git a/examples/molerat/downsample_source_images.py b/examples/molerat/1_downsample_source_images.py similarity index 77% rename from examples/molerat/downsample_source_images.py rename to examples/molerat/1_downsample_source_images.py index 0e58462..7ae04e4 100644 --- a/examples/molerat/downsample_source_images.py +++ b/examples/molerat/1_downsample_source_images.py @@ -8,6 +8,8 @@ from loguru import logger from skimage import transform +from brainglobe_template_builder.plots import plot_grid, plot_orthographic + if __name__ == "__main__": parser = argparse.ArgumentParser(description="Downsample source images") parser.add_argument( @@ -75,7 +77,7 @@ rawdata_filename = ( f"sub-{subject_id}_" f"hemi-{hemisphere}_" - f"res-{target_isotropic_resolution}.tif" + f"res-{target_isotropic_resolution}um.tif" ) # we can't use our usual transform utils function here, @@ -88,14 +90,36 @@ sample["image_orientation"] == "RPI" ), "Image orientation is not RPI" stack = load_any(str(source_file)) + # Find the last few slices of stack that contain all zeros + zero_slices = 0 + for i in range(stack.shape[0] - 1, -1, -1): + if np.all(stack[i] == 0): + zero_slices += 1 + else: + break + if zero_slices: + logger.info(f"Last zero slices: {zero_slices}") + stack = stack[: stack.shape[0] - zero_slices, :, :] + + # mirror BEFORE downsampling + # (because otherwise midline gets blurred by zeros!) + stack = np.concatenate((stack, np.flip(stack, axis=0)), axis=0) downsampled = transform.downscale_local_mean( stack, (axial_factor, in_plane_factor, in_plane_factor) ) - downsampled = np.concatenate( - (downsampled, np.flip(downsampled, axis=0)), axis=0 - ) original_space = AnatomicalSpace("RPI") downsampled = original_space.map_stack_to("ASR", downsampled) save_any(downsampled, subject_folder / rawdata_filename) + plots_folder = ( + Path.home() / "dev/brainglobe-template-builder/test-images/" + ) + plot_grid( + downsampled, + save_path=plots_folder / f"grid-{rawdata_filename}.png", + ) + plot_orthographic( + downsampled, + save_path=plots_folder / f"ortho-{rawdata_filename}.png", + ) logger.info(f"{rawdata_filename} downsampled.") diff --git a/examples/molerat/2_molerat_prep_lowres.py b/examples/molerat/2_molerat_prep_lowres.py new file mode 100644 index 0000000..8aea0d4 --- /dev/null +++ b/examples/molerat/2_molerat_prep_lowres.py @@ -0,0 +1,192 @@ +""" +Prepare low-resolution Molerat images for template construction +================================================================ +The following operations are performed on the lowest-resolution images to +prepare them for template construction: +- Perform N4 Bias field correction using ANTs +- Generate brain mask based on N4-corrected image +- Save all resulting images as nifti files to be used for template construction +""" + +# %% +# Imports +# ------- +import os +from datetime import date +from pathlib import Path + +import ants +import numpy as np +from loguru import logger + +from brainglobe_template_builder.io import ( + file_path_with_suffix, + load_tiff, + save_as_asr_nii, +) +from brainglobe_template_builder.preproc.masking import create_mask +from brainglobe_template_builder.preproc.splitting import ( + save_array_dict_to_nii, +) + +# %% +# Setup +# ----- +# Define some global variables, including paths to input and output directories +# and set up logging. + +# Define voxel size(in microns) of the lowest resolution image +lowres = 40 +# String to identify the resolution in filenames +res_str = f"res-{lowres}um" +# Define voxel sizes in mm (for Nifti saving) +vox_sizes = [lowres * 1e-3] * 3 # in mm + +# Prepare directory structure +atlas_dir = Path("/media/ceph-neuroinformatics/neuroinformatics/atlas-forge") +species_id = "MoleRat" +species_dir = atlas_dir / species_id +raw_dir = species_dir / "rawdata" +assert raw_dir.exists(), f"Could not find rawdata directory {raw_dir}." +deriv_dir = species_dir / "derivatives" +assert deriv_dir.exists(), f"Could not find derivatives directory {deriv_dir}." + +# Set up logging +today = date.today() +current_script_name = os.path.basename(__file__).replace(".py", "") +logger.add(species_dir / "logs" / f"{today}_{current_script_name}.log") + + +# %% +# Run the pipeline for each subject +# --------------------------------- + +# Create a dictionary to store the paths to the use4template directories +# per subject. These will contain all necessary images for template building. +use4template_dirs = {} + +for raw_subj_dir in raw_dir.iterdir(): + + file_prefix = f"{raw_subj_dir.name}_{res_str}" + deriv_subj_dir = deriv_dir / raw_subj_dir.name + deriv_subj_dir.mkdir(exist_ok=True) + raw_tiff_path = raw_subj_dir / f"{file_prefix}.tif" + assert raw_tiff_path.exists() + + logger.info(f"Starting to process {file_prefix}...") + logger.info(f"Will save outputs to {deriv_subj_dir}/") + + # Load the image (already in ASR) + image = load_tiff(raw_tiff_path) + logger.debug( + f"Loaded image {raw_tiff_path.name} with shape: {image.shape}." + ) + + # Save the image as nifti + nii_path = file_path_with_suffix( + deriv_subj_dir / f"{file_prefix}.tif", "_orig-asr", new_ext=".nii.gz" + ) + save_as_asr_nii(image, vox_sizes, nii_path) + logger.debug(f"Saved reoriented image as {nii_path.name}.") + + # Bias field correction (to homogenise intensities) + image_ants = ants.image_read(nii_path.as_posix()) + image_n4 = ants.n4_bias_field_correction(image_ants) + image_n4_path = file_path_with_suffix(nii_path, "_N4") + ants.image_write(image_n4, image_n4_path.as_posix()) + logger.debug( + f"Created N4 bias field corrected image as {image_n4_path.name}." + ) + + # Generate a brain mask based on the N4-corrected image + mask_data = create_mask( + image_n4.numpy(), + gauss_sigma=3, + threshold_method="triangle", + closing_size=5, + ) + mask_path = file_path_with_suffix(nii_path, "_N4_mask") + mask = image_n4.new_image_like(mask_data.astype(np.uint8)) + ants.image_write(mask, mask_path.as_posix()) + logger.debug( + f"Generated brain mask with shape: {mask.shape} " + f"and saved as {mask_path.name}." + ) + + # Plot the mask over the image to check + mask_plot_path = ( + deriv_subj_dir / f"{file_prefix}_orig-asr_N4_mask-overlay.png" + ) + ants.plot( + image_n4, + mask, + overlay_alpha=0.5, + axis=1, + title="Brain mask over image", + filename=mask_plot_path.as_posix(), + ) + logger.debug("Plotted overlay to visually check mask.") + + output_prefix = file_path_with_suffix(nii_path, "_N4_aligned_", new_ext="") + # Generate arrays for template construction and save as niftis + use4template_dir = Path(output_prefix.as_posix() + "padded_use4template") + # if it exists, delete existing files in it + if use4template_dir.exists(): + logger.warning(f"Removing existing files in {use4template_dir}.") + for file in use4template_dir.glob("*"): + file.unlink() + use4template_dir.mkdir(exist_ok=True) + + array_dict = { + f"{use4template_dir}/{file_prefix}_sym-brain": np.pad( + image_n4.numpy(), pad_width=2, mode="constant" + ), + f"{use4template_dir}/{file_prefix}_sym-mask": np.pad( + mask.numpy(), pad_width=2, mode="constant" + ), + } + save_array_dict_to_nii(array_dict, use4template_dir, vox_sizes) + use4template_dirs[file_prefix] = use4template_dir + logger.info( + f"Saved images for template construction in {use4template_dir}." + ) + logger.info(f"Finished processing {file_prefix}.") + +# %% +# Generate lists of file paths for template construction +# ----------------------------------------------------- +# Use the paths to the use4template directories to generate lists of file paths +# for the template construction pipeline. Three kinds of template will be +# generated, and each needs the corresponging brain image and mask files: +# 1. All asym* files for subjects where hemi=both. These will be used to +# generate an asymmetric brain template. +# 2. All right-sym* files for subjects where use_right is True, and +# all left-sym* files for subjects where use_left is True. +# These will be used to generate a symmetric brain template. +# 3. All right-hemi* files for subjects where use_right is True, +# and all left-hemi-xflip* files for subjects where use_left is True. +# These will be used to generate a symmetric hemisphere template. + +filepath_lists: dict[str, list] = { + "sym-brain": [], + "sym-mask": [], +} + +for subject, use4template_dir in use4template_dirs.items(): + for label in ["brain", "mask"]: + filepath_lists[f"sym-{label}"].append( + use4template_dir / f"{subject}_sym-{label}.nii.gz" + ) + +# %% +# Save the file paths to text files, each in a separate directory + +for key, paths in filepath_lists.items(): + kind, label = key.split("-") # e.g. "asym" and "brain" + n_images = len(paths) + template_name = f"template_{kind}_{res_str}_n-{n_images}" + template_dir = species_dir / "templates" / template_name + template_dir.mkdir(exist_ok=True) + np.savetxt(template_dir / f"{label}_paths.txt", paths, fmt="%s") + +# %% diff --git a/examples/slurm/build_template_molerat.sh b/examples/slurm/build_template_molerat.sh new file mode 100644 index 0000000..0adee74 --- /dev/null +++ b/examples/slurm/build_template_molerat.sh @@ -0,0 +1,30 @@ +#!/bin/bash + +#SBATCH -J temp_build # job name +#SBATCH -p cpu # partition +#SBATCH -N 1 # number of nodes +#SBATCH --mem 64G # memory pool for all cores +#SBATCH -n 20 # number of cores +#SBATCH -t 5-00:00 # time (D-HH:MM) +#SBATCH -o slurm.%x.%N.%j.out # write STDOUT +#SBATCH -e slurm.%x.%N.%j.err # write STDERR +#SBATCH --mail-type=ALL +#SBATCH --mail-user=cceafel@ucl.ac.uk + +# Load the required modules +module load miniconda +module load ants +# Activate the conda environment +mamba activate template-builder + +# Define atlas-forge directory, species and template names, and average type +ATLAS_DIR="/ceph/neuroinformatics/neuroinformatics/atlas-forge" +SPECIES="MoleRat" +TEMP_NAME="template_sym_res-40um_n-6" +AVE_TYPE="trimmed_mean" + +# Path to the bash script that builds the template +BUILD_SCRIPT="${ATLAS_DIR}/${SPECIES}/scripts/MoleRat_build_template.sh" + +# Run the script to build the template +bash $BUILD_SCRIPT --atlas-dir $ATLAS_DIR --species $SPECIES --template-name $TEMP_NAME --average-type $AVE_TYPE From 8478d073b66c6fd37cbe957d4cd06078975d5f22 Mon Sep 17 00:00:00 2001 From: alessandrofelder Date: Fri, 29 Nov 2024 11:14:44 +0000 Subject: [PATCH 3/3] initial molerat template build script --- examples/molerat/3_build_template.sh | 99 ++++++++++++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100755 examples/molerat/3_build_template.sh diff --git a/examples/molerat/3_build_template.sh b/examples/molerat/3_build_template.sh new file mode 100755 index 0000000..b732c80 --- /dev/null +++ b/examples/molerat/3_build_template.sh @@ -0,0 +1,99 @@ +#!/bin/bash + +# Start the timer +start_time=$(date +%s) + +# Initialize CLI parameters with default values +average_type="mean" + +# Function to display help message +usage() { + echo "Usage: $0 --atlas-dir --template-name [--average-type ]" + echo "" + echo "Options:" + echo " --atlas-dir Path to the atlas-forge directory [REQUIRED]" + echo " --template-name The name of the appropriate subfolder within templates (e.g., template_asym_res-50um_n-8) [REQUIRED]" + echo " --average-type The type of average to use (default: mean)." + exit 1 +} + +# Check for help flag first +if [[ "$1" == "--help" ]]; then + usage +fi + +# Parse command line arguments +while [[ $# -gt 0 ]]; do + case $1 in + --atlas-dir) + atlas_dir="$2" + shift 2 + ;; + --template-name) + template_name="$2" + shift 2 + ;; + --average-type) + average_type="$2" + shift 2 + ;; + --help) + usage + ;; + *) + echo "Unknown option $1" + usage + ;; + esac +done + +# Check for required parameters +if [ -z "$atlas_dir" ] || [ -z "$template_name" ]; then + echo "Error: --atlas-dir and --template-name are required." + usage +fi + +echo "atlas-dir: ${atlas_dir}" +echo "Building the template ${template_name}..." + +# Set the path to the working directory where the template will be built +templates_dir="${atlas_dir}/templates" +working_dir="${templates_dir}/${template_name}" + +# If average type is trimmed_mean or efficient_trimean, we need python +average_prog="ANTs" +if [ "${average_type}" == "trimmed_mean" ] || [ "${average_type}" == "efficient_trimean" ]; then + average_prog="python" +fi + +# Verify that the working directory exists before changing directory +if [ ! -d "${working_dir}" ]; then + echo "The working directory does not exist: ${working_dir}" + exit 1 +fi + +cd "${working_dir}" || exit + +echo "Results will be written to ${working_dir}" + +# Execute the actual template building +echo "Starting to build the template..." +bash modelbuild.sh --output-dir "${working_dir}" \ + --starting-target first \ + --stages rigid,similarity,affine,nlin \ + --masks "${working_dir}/mask_paths.txt" \ + --average-type "${average_type}" \ + --average-prog "${average_prog}" \ + --reuse-affines \ + --dry-run \ + "${working_dir}/brain_paths.txt" + +echo "Finished building the template!" + +# Write execution time (in HH:MM format) to a file +end_time=$(date +%s) +execution_time=$((end_time - start_time)) +hours=$((execution_time / 3600)) +minutes=$(( (execution_time % 3600) / 60)) +formatted_time=$(printf "%02d:%02d" $hours $minutes) +echo "Execution time: $formatted_time" > "${working_dir}/execution_time.txt"