From 4b16d8626694b0a2317a1adc18ca8a642431e2f4 Mon Sep 17 00:00:00 2001
From: niksirbi <niko.sirbiladze@gmail.com>
Date: Mon, 19 Aug 2024 18:40:27 +0100
Subject: [PATCH 1/9] wip adapt prep_lowres for tadpoles

---
 examples/tadpole/prep_lowres.py | 337 ++++++++++++++++++++++++++++++++
 1 file changed, 337 insertions(+)
 create mode 100644 examples/tadpole/prep_lowres.py

diff --git a/examples/tadpole/prep_lowres.py b/examples/tadpole/prep_lowres.py
new file mode 100644
index 0000000..d507001
--- /dev/null
+++ b/examples/tadpole/prep_lowres.py
@@ -0,0 +1,337 @@
+"""
+Prepare low-resolution BlackCap images for template construction
+================================================================
+The following operations are performed on the lowest-resolution images to
+prepare them for template construction:
+- Import each image as tiff, re-orient to ASR and save as nifti
+- Perform N4 Bias field correction using ANTs
+- Generate brain mask based on N4-corrected image
+- Rigid-register the re-oriented image to an already aligned target (with ANTs)
+- Transform the image and mask into the aligned space
+- Split the image and mask into hemispheres and reflect each hemisphere
+- Generate symmetric brains using either the left or right hemisphere
+- Save all resulting images as nifti files to be used for template construction
+"""
+
+# %%
+# Imports
+# -------
+import os
+import subprocess
+from datetime import date
+from pathlib import Path
+
+import ants
+import numpy as np
+import pandas as pd
+from brainglobe_space import AnatomicalSpace
+from loguru import logger
+from tqdm import tqdm
+
+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 (
+    generate_arrays_4template,
+    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 = 12
+# 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
+
+# Masking parameters
+mask_params = {
+    "gauss_sigma": 2,
+    "threshold_method": "triangle",
+    "closing_size": 3,
+}
+
+# Prepare directory structure
+atlas_forge_dir = Path("/media/ceph-niu/neuroinformatics/atlas-forge")
+project_dir = (
+    atlas_forge_dir / "Tadpole" / "tadpole-template-starter" / "young-tadpole"
+)
+raw_dir = project_dir / "rawdata"
+deriv_dir = project_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(project_dir / "logs" / f"{today}_{current_script_name}.log")
+
+# %%
+# Define registration target
+# --------------------------
+# We need three images, all in ASR orientation and nifti format:
+# 1. An initial target aligned to the mid-sagittal plane
+# 2. The brain mask of the initial target image
+# 3. The mask of the halves of the initial target image (left and right)
+
+# Here we used one of the young tadpole images (sub-najva4) as initial target.
+# We have manually aligned the target to the mid-sagittal plane with napari
+# and prepared the masks accordingly.
+
+target_dir = project_dir / "templates" / "initial-target"
+target_prefix = f"sub-najva4_{res_str}_channel-orange_orig-asr"
+target_image_path = target_dir / f"{target_prefix}_aligned.nii.gz"
+target_image = ants.image_read(target_image_path.as_posix())
+target_mask_path = target_dir / f"{target_prefix}_label-brain_aligned.nii.gz"
+target_mask = ants.image_read(target_mask_path.as_posix())
+target_halves_mask = ants.image_read(
+    (target_dir / f"{target_prefix}_label-halves_aligned.nii.gz").as_posix()
+)
+
+# %%
+# Load a dataframe with image paths to use for the template
+# ---------------------------------------------------------
+
+source_csv_dir = project_dir / "templates" / "use_for_template.csv"
+df = pd.read_csv(source_csv_dir)
+
+# take the hemi column and split it into two boolean columns
+# named use_left and use_right (for easier filtering)
+df["use_left"] = df["hemi"].isin(["left", "both"])
+df["use_right"] = df["hemi"].isin(["right", "both"])
+# Keep only the rows in which at least one of use_left or use_right is True
+df = df[df[["use_left", "use_right"]].any(axis=1)]
+n_subjects = len(df)
+
+# %%
+# 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 idx, row in tqdm(df.iterrows(), total=n_subjects):
+    # Figure out input-output paths
+    row = df.iloc[idx]
+    subject = "sub-" + row["subject_id"]
+    channel_str = "channel-" + row["color"]
+    source_origin = row["orientation"]
+    file_prefix = f"{subject}_{res_str}_{channel_str}"
+    deriv_subj_dir = deriv_dir / subject
+    deriv_subj_dir.mkdir(exist_ok=True)
+    tiff_path = raw_dir / f"{file_prefix}.tif"
+
+    logger.info(f"Starting to process {file_prefix}...")
+    logger.info(f"Will save outputs to {deriv_subj_dir}/")
+
+    # Load the image
+    image = load_tiff(tiff_path)
+    logger.debug(f"Loaded image {tiff_path.name} with shape: {image.shape}.")
+
+    # Reorient the image to ASR
+    target_origin = "asr"
+    source_space = AnatomicalSpace(source_origin, shape=image.shape)
+    image_asr = source_space.map_stack_to(target_origin, image)
+    logger.debug(f"Reoriented image from {source_origin} to {target_origin}.")
+    logger.debug(f"Reoriented image shape: {image_asr.shape}.")
+
+    # Save the reoriented image as nifti
+    nii_path = file_path_with_suffix(
+        deriv_subj_dir / tiff_path.name, f"_{target_origin}", new_ext=".nii.gz"
+    )
+    save_as_asr_nii(image_asr, 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(), **mask_params)  # type: ignore
+    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.")
+
+    # Rigid-register the reoriented image to an already aligned target
+    output_prefix = file_path_with_suffix(nii_path, "_N4_aligned_", new_ext="")
+
+    # Call the antsRegistration_affine_SyN.sh script with provided parameters
+    cmd = "antsRegistration_affine_SyN.sh "
+    cmd += f"--moving-mask {mask_path.as_posix()} "
+    cmd += f"--fixed-mask {target_mask_path.as_posix()} "
+    cmd += "--linear-type rigid "  # Use rigid registration
+    cmd += "--skip-nonlinear "  # Skip non-linear registration
+    cmd += "--clobber "  # Overwrite existing files
+    cmd += f"{image_n4_path.as_posix()} "  # moving image
+    cmd += f"{target_image_path.as_posix()} "  # fixed image
+    cmd += f"{output_prefix.as_posix()}"  # output prefix
+
+    logger.debug(f"Running the following ANTs registration script: {cmd}")
+    subprocess.run(cmd, shell=True, check=True)
+    logger.debug(
+        "Computed rigid transformation to align image to initial target."
+    )
+
+    # Load the computed rigid transformation
+    rigid_transform_path = file_path_with_suffix(
+        output_prefix, "0GenericAffine.mat"
+    )
+    assert (
+        rigid_transform_path.exists()
+    ), f"Could not find the rigid transformation at {rigid_transform_path}."
+    # Use the rigid transformation to align the image
+    aligned_image = ants.apply_transforms(
+        fixed=target_image,
+        moving=image_n4,
+        transformlist=rigid_transform_path.as_posix(),
+        interpolator="bSpline",
+    )
+    logger.debug("Transformed image to aligned space.")
+    # Also align the mask
+    aligned_mask = ants.apply_transforms(
+        fixed=target_image,
+        moving=mask,
+        transformlist=rigid_transform_path.as_posix(),
+        interpolator="nearestNeighbor",
+    )
+    logger.debug("Transformed brain mask to aligned space.")
+
+    # Plot the aligned image over the target to check registration
+    ants.plot(
+        target_image,
+        aligned_image,
+        overlay_alpha=0.5,
+        overlay_cmap="plasma",
+        axis=1,
+        title="Aligned image over target (rigid registration)",
+        filename=output_prefix.as_posix() + "target-overlay.png",
+    )
+    # Plot the halves mask over the aligned image to check the split
+    ants.plot(
+        aligned_image,
+        target_halves_mask,
+        overlay_alpha=0.5,
+        axis=1,
+        title="Aligned image split into right and left halves",
+        filename=output_prefix.as_posix() + "halves-overlay.png",
+    )
+    logger.debug("Plotted overlays to visually check alignment.")
+
+    # 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 = generate_arrays_4template(
+        subject, aligned_image.numpy(), aligned_mask.numpy(), pad=2
+    )
+    save_array_dict_to_nii(array_dict, use4template_dir, vox_sizes)
+    use4template_dirs[subject] = 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] = {
+    "asym-brain": [],
+    "asym-mask": [],
+    "sym-brain": [],
+    "sym-mask": [],
+    "hemi-brain": [],
+    "hemi-mask": [],
+}
+
+for _, row in df.iterrows():
+    subject = "sub-" + row["subject_id"]
+    use4template_dir = use4template_dirs[subject]
+
+    if row["hemi"] == "both":
+        # Add paths for the asymmetric brain template
+        for label in ["brain", "mask"]:
+            filepath_lists[f"asym-{label}"].append(
+                use4template_dir / f"{subject}_asym-{label}.nii.gz"
+            )
+
+    if row["use_right"]:
+        for label in ["brain", "mask"]:
+            # Add paths for the symmetric brain template
+            filepath_lists[f"sym-{label}"].append(
+                use4template_dir / f"{subject}_right-sym-{label}.nii.gz"
+            )
+            # Add paths for the hemispheric template
+            filepath_lists[f"hemi-{label}"].append(
+                use4template_dir / f"{subject}_right-hemi-{label}.nii.gz"
+            )
+
+    if row["use_left"]:
+        for label in ["brain", "mask"]:
+            # Add paths for the symmetric brain template
+            filepath_lists[f"sym-{label}"].append(
+                use4template_dir / f"{subject}_left-sym-{label}.nii.gz"
+            )
+            # Add paths for the hemispheric template
+            filepath_lists[f"hemi-{label}"].append(
+                use4template_dir / f"{subject}_left-hemi-xflip-{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 = project_dir / "templates" / template_name
+    template_dir.mkdir(exist_ok=True)
+    np.savetxt(template_dir / f"{label}_paths.txt", paths, fmt="%s")

From d1792dbffa2d7a6ec0c2e5f0031198e4db209744 Mon Sep 17 00:00:00 2001
From: alessandrofelder <alessandrofelder@users.noreply.github.com>
Date: Wed, 21 Aug 2024 15:47:57 +0100
Subject: [PATCH 2/9] adapt tadpole prep lowres script for manually aligned
 data

---
 examples/tadpole/prep_lowres.py | 101 +++++---------------------------
 1 file changed, 16 insertions(+), 85 deletions(-)

diff --git a/examples/tadpole/prep_lowres.py b/examples/tadpole/prep_lowres.py
index d507001..35c8f76 100644
--- a/examples/tadpole/prep_lowres.py
+++ b/examples/tadpole/prep_lowres.py
@@ -17,23 +17,18 @@
 # Imports
 # -------
 import os
-import subprocess
 from datetime import date
 from pathlib import Path
 
 import ants
 import numpy as np
 import pandas as pd
-from brainglobe_space import AnatomicalSpace
 from loguru import logger
 from tqdm import tqdm
 
 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 (
     generate_arrays_4template,
     save_array_dict_to_nii,
@@ -62,7 +57,7 @@
 # Prepare directory structure
 atlas_forge_dir = Path("/media/ceph-niu/neuroinformatics/atlas-forge")
 project_dir = (
-    atlas_forge_dir / "Tadpole" / "tadpole-template-starter" / "young-tadpole"
+    atlas_forge_dir / "Tadpole" / "tadpole-template-starter" / "old-tadpole"
 )
 raw_dir = project_dir / "rawdata"
 deriv_dir = project_dir / "derivatives"
@@ -86,7 +81,7 @@
 # and prepared the masks accordingly.
 
 target_dir = project_dir / "templates" / "initial-target"
-target_prefix = f"sub-najva4_{res_str}_channel-orange_orig-asr"
+target_prefix = f"sub-topro54_{res_str}_channel-orange_orig-asr"
 target_image_path = target_dir / f"{target_prefix}_aligned.nii.gz"
 target_image = ants.image_read(target_image_path.as_posix())
 target_mask_path = target_dir / f"{target_prefix}_label-brain_aligned.nii.gz"
@@ -132,23 +127,10 @@
     logger.info(f"Starting to process {file_prefix}...")
     logger.info(f"Will save outputs to {deriv_subj_dir}/")
 
-    # Load the image
-    image = load_tiff(tiff_path)
-    logger.debug(f"Loaded image {tiff_path.name} with shape: {image.shape}.")
-
-    # Reorient the image to ASR
-    target_origin = "asr"
-    source_space = AnatomicalSpace(source_origin, shape=image.shape)
-    image_asr = source_space.map_stack_to(target_origin, image)
-    logger.debug(f"Reoriented image from {source_origin} to {target_origin}.")
-    logger.debug(f"Reoriented image shape: {image_asr.shape}.")
-
-    # Save the reoriented image as nifti
+    # Load the manually aligned nifti
     nii_path = file_path_with_suffix(
-        deriv_subj_dir / tiff_path.name, f"_{target_origin}", new_ext=".nii.gz"
+        deriv_subj_dir / tiff_path.name, "_orig-asr_aligned", new_ext=".nii.gz"
     )
-    save_as_asr_nii(image_asr, 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())
@@ -159,14 +141,16 @@
         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(), **mask_params)  # type: ignore
-    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())
+    # Read the manually adjusted brain mask
+    # Save the reoriented image as nifti
+    mask_path = file_path_with_suffix(
+        deriv_subj_dir / tiff_path.name,
+        "_orig-asr_label-brain_aligned",
+        new_ext=".nii.gz",
+    )
+    mask = ants.image_read(mask_path.as_posix())
     logger.debug(
-        f"Generated brain mask with shape: {mask.shape} "
-        f"and saved as {mask_path.name}."
+        f"Read brain mask with shape: {mask.shape} " f"from {mask_path.name}."
     )
 
     # Plot the mask over the image to check
@@ -183,63 +167,10 @@
     )
     logger.debug("Plotted overlay to visually check mask.")
 
-    # Rigid-register the reoriented image to an already aligned target
-    output_prefix = file_path_with_suffix(nii_path, "_N4_aligned_", new_ext="")
-
-    # Call the antsRegistration_affine_SyN.sh script with provided parameters
-    cmd = "antsRegistration_affine_SyN.sh "
-    cmd += f"--moving-mask {mask_path.as_posix()} "
-    cmd += f"--fixed-mask {target_mask_path.as_posix()} "
-    cmd += "--linear-type rigid "  # Use rigid registration
-    cmd += "--skip-nonlinear "  # Skip non-linear registration
-    cmd += "--clobber "  # Overwrite existing files
-    cmd += f"{image_n4_path.as_posix()} "  # moving image
-    cmd += f"{target_image_path.as_posix()} "  # fixed image
-    cmd += f"{output_prefix.as_posix()}"  # output prefix
-
-    logger.debug(f"Running the following ANTs registration script: {cmd}")
-    subprocess.run(cmd, shell=True, check=True)
-    logger.debug(
-        "Computed rigid transformation to align image to initial target."
-    )
-
-    # Load the computed rigid transformation
-    rigid_transform_path = file_path_with_suffix(
-        output_prefix, "0GenericAffine.mat"
-    )
-    assert (
-        rigid_transform_path.exists()
-    ), f"Could not find the rigid transformation at {rigid_transform_path}."
-    # Use the rigid transformation to align the image
-    aligned_image = ants.apply_transforms(
-        fixed=target_image,
-        moving=image_n4,
-        transformlist=rigid_transform_path.as_posix(),
-        interpolator="bSpline",
-    )
-    logger.debug("Transformed image to aligned space.")
-    # Also align the mask
-    aligned_mask = ants.apply_transforms(
-        fixed=target_image,
-        moving=mask,
-        transformlist=rigid_transform_path.as_posix(),
-        interpolator="nearestNeighbor",
-    )
-    logger.debug("Transformed brain mask to aligned space.")
-
-    # Plot the aligned image over the target to check registration
-    ants.plot(
-        target_image,
-        aligned_image,
-        overlay_alpha=0.5,
-        overlay_cmap="plasma",
-        axis=1,
-        title="Aligned image over target (rigid registration)",
-        filename=output_prefix.as_posix() + "target-overlay.png",
-    )
     # Plot the halves mask over the aligned image to check the split
+    output_prefix = file_path_with_suffix(nii_path, "_N4_aligned_", new_ext="")
     ants.plot(
-        aligned_image,
+        image_n4,
         target_halves_mask,
         overlay_alpha=0.5,
         axis=1,
@@ -258,7 +189,7 @@
     use4template_dir.mkdir(exist_ok=True)
 
     array_dict = generate_arrays_4template(
-        subject, aligned_image.numpy(), aligned_mask.numpy(), pad=2
+        subject, image_n4.numpy(), mask.numpy(), pad=2
     )
     save_array_dict_to_nii(array_dict, use4template_dir, vox_sizes)
     use4template_dirs[subject] = use4template_dir

From 8ace1782660ff39e50f932699f6a8eba733eecdb Mon Sep 17 00:00:00 2001
From: alessandrofelder <alessandrofelder@users.noreply.github.com>
Date: Wed, 21 Aug 2024 16:19:16 +0100
Subject: [PATCH 3/9] draft build_template.sh

---
 examples/tadpole/build_template.sh | 99 ++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)
 create mode 100755 examples/tadpole/build_template.sh

diff --git a/examples/tadpole/build_template.sh b/examples/tadpole/build_template.sh
new file mode 100755
index 0000000..b732c80
--- /dev/null
+++ b/examples/tadpole/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 <path> --template-name <string> [--average-type <string>]"
+  echo ""
+  echo "Options:"
+  echo "  --atlas-dir <path>         Path to the atlas-forge directory [REQUIRED]"
+  echo "  --template-name <string>   The name of the appropriate subfolder within templates (e.g., template_asym_res-50um_n-8) [REQUIRED]"
+  echo "  --average-type <string>    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"

From bb5c040726b0a7e0a213cbc9bf611e19ef3aaebe Mon Sep 17 00:00:00 2001
From: alessandrofelder <alessandrofelder@users.noreply.github.com>
Date: Thu, 22 Aug 2024 15:49:21 +0100
Subject: [PATCH 4/9] additionally apply mask when N4 correcting

---
 examples/tadpole/prep_lowres.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/examples/tadpole/prep_lowres.py b/examples/tadpole/prep_lowres.py
index 35c8f76..6d6c644 100644
--- a/examples/tadpole/prep_lowres.py
+++ b/examples/tadpole/prep_lowres.py
@@ -132,15 +132,6 @@
         deriv_subj_dir / tiff_path.name, "_orig-asr_aligned", new_ext=".nii.gz"
     )
 
-    # 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}."
-    )
-
     # Read the manually adjusted brain mask
     # Save the reoriented image as nifti
     mask_path = file_path_with_suffix(
@@ -153,6 +144,17 @@
         f"Read brain mask with shape: {mask.shape} " f"from {mask_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_masked_numpy = image_n4.numpy() * mask.numpy()
+    image_n4_path = file_path_with_suffix(nii_path, "_N4")
+    image_n4 = image_n4.new_image_like(image_n4_masked_numpy)
+    ants.image_write(image_n4, image_n4_path.as_posix())
+    logger.debug(
+        f"Created N4 bias field corrected image as {image_n4_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"

From 392ea70e662f6c9e8c65b94d3012791a542dbb72 Mon Sep 17 00:00:00 2001
From: niksirbi <niko.sirbiladze@gmail.com>
Date: Fri, 18 Oct 2024 14:38:01 +0100
Subject: [PATCH 5/9] renamed and cleaned up prep-lowres script for tadpoles

---
 .../{prep_lowres.py => 2_prep_lowres.py}      | 73 +++++++++----------
 1 file changed, 35 insertions(+), 38 deletions(-)
 rename examples/tadpole/{prep_lowres.py => 2_prep_lowres.py} (82%)

diff --git a/examples/tadpole/prep_lowres.py b/examples/tadpole/2_prep_lowres.py
similarity index 82%
rename from examples/tadpole/prep_lowres.py
rename to examples/tadpole/2_prep_lowres.py
index 6d6c644..da8707d 100644
--- a/examples/tadpole/prep_lowres.py
+++ b/examples/tadpole/2_prep_lowres.py
@@ -1,13 +1,16 @@
 """
 Prepare low-resolution BlackCap images for template construction
 ================================================================
-The following operations are performed on the lowest-resolution images to
-prepare them for template construction:
-- Import each image as tiff, re-orient to ASR and save as nifti
-- Perform N4 Bias field correction using ANTs
-- Generate brain mask based on N4-corrected image
-- Rigid-register the re-oriented image to an already aligned target (with ANTs)
-- Transform the image and mask into the aligned space
+The ``brainglobe-template-builder`` Preprocess widget was used to perform
+the following operations on each subject's low-resolution image, ahead of
+running this script:
+- Load each image from tiff, re-orient to ASR and save as nifti
+- Generate a brain mask automatically and modify it manually, if necessary
+- Find the mid-sagittal plane and align the image and mask accordingly
+- Save the aligned image, brain mask, and halves mask as nifti files
+
+This script will:
+- Perform N4 Bias field correction on the aligned image using ANTs
 - Split the image and mask into hemispheres and reflect each hemisphere
 - Generate symmetric brains using either the left or right hemisphere
 - Save all resulting images as nifti files to be used for template construction
@@ -68,28 +71,6 @@
 current_script_name = os.path.basename(__file__).replace(".py", "")
 logger.add(project_dir / "logs" / f"{today}_{current_script_name}.log")
 
-# %%
-# Define registration target
-# --------------------------
-# We need three images, all in ASR orientation and nifti format:
-# 1. An initial target aligned to the mid-sagittal plane
-# 2. The brain mask of the initial target image
-# 3. The mask of the halves of the initial target image (left and right)
-
-# Here we used one of the young tadpole images (sub-najva4) as initial target.
-# We have manually aligned the target to the mid-sagittal plane with napari
-# and prepared the masks accordingly.
-
-target_dir = project_dir / "templates" / "initial-target"
-target_prefix = f"sub-topro54_{res_str}_channel-orange_orig-asr"
-target_image_path = target_dir / f"{target_prefix}_aligned.nii.gz"
-target_image = ants.image_read(target_image_path.as_posix())
-target_mask_path = target_dir / f"{target_prefix}_label-brain_aligned.nii.gz"
-target_mask = ants.image_read(target_mask_path.as_posix())
-target_halves_mask = ants.image_read(
-    (target_dir / f"{target_prefix}_label-halves_aligned.nii.gz").as_posix()
-)
-
 # %%
 # Load a dataframe with image paths to use for the template
 # ---------------------------------------------------------
@@ -122,30 +103,46 @@
     file_prefix = f"{subject}_{res_str}_{channel_str}"
     deriv_subj_dir = deriv_dir / subject
     deriv_subj_dir.mkdir(exist_ok=True)
-    tiff_path = raw_dir / f"{file_prefix}.tif"
+    tiff_path = raw_dir / f"{file_prefix}.tif"  # original tiff image
+    nii_ext = ".nii.gz"  # extension for nifti files
 
     logger.info(f"Starting to process {file_prefix}...")
     logger.info(f"Will save outputs to {deriv_subj_dir}/")
 
-    # Load the manually aligned nifti
+    # Load the manually aligned nifti image
     nii_path = file_path_with_suffix(
-        deriv_subj_dir / tiff_path.name, "_orig-asr_aligned", new_ext=".nii.gz"
+        deriv_subj_dir / tiff_path.name, "_orig-asr_aligned", new_ext=nii_ext
+    )
+    image_ants = ants.image_read(nii_path.as_posix())
+    logger.debug(
+        f"Read image with shape: {image_ants.shape} from {nii_path.name}."
     )
 
-    # Read the manually adjusted brain mask
-    # Save the reoriented image as nifti
+    # Load the manually adjusted brain mask for the image
     mask_path = file_path_with_suffix(
         deriv_subj_dir / tiff_path.name,
         "_orig-asr_label-brain_aligned",
-        new_ext=".nii.gz",
+        new_ext=nii_ext,
     )
     mask = ants.image_read(mask_path.as_posix())
     logger.debug(
-        f"Read brain mask with shape: {mask.shape} " f"from {mask_path.name}."
+        f"Read brain mask with shape: {mask.shape} from {mask_path.name}."
+    )
+
+    # Load the halves mask for the image (also pre-generated via the widget)
+    # This will be later used for diagnostic plotting
+    halves_mask_path = file_path_with_suffix(
+        deriv_subj_dir / tiff_path.name,
+        "_orig-asr_label-halves_aligned",
+        new_ext=nii_ext,
+    )
+    halves_mask = ants.image_read(halves_mask_path.as_posix())
+    logger.debug(
+        f"Read halves mask with shape: {halves_mask.shape} from "
+        f"{halves_mask_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_masked_numpy = image_n4.numpy() * mask.numpy()
     image_n4_path = file_path_with_suffix(nii_path, "_N4")
@@ -173,7 +170,7 @@
     output_prefix = file_path_with_suffix(nii_path, "_N4_aligned_", new_ext="")
     ants.plot(
         image_n4,
-        target_halves_mask,
+        halves_mask,
         overlay_alpha=0.5,
         axis=1,
         title="Aligned image split into right and left halves",

From 654b6ad0a806aebcbd62e2e3786f18983b1c68d0 Mon Sep 17 00:00:00 2001
From: niksirbi <niko.sirbiladze@gmail.com>
Date: Fri, 18 Oct 2024 14:40:40 +0100
Subject: [PATCH 6/9] renamed tadpole build template script

---
 examples/tadpole/{build_template.sh => 3_build_template.sh} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename examples/tadpole/{build_template.sh => 3_build_template.sh} (100%)

diff --git a/examples/tadpole/build_template.sh b/examples/tadpole/3_build_template.sh
similarity index 100%
rename from examples/tadpole/build_template.sh
rename to examples/tadpole/3_build_template.sh

From f163a3ac584fbbe4da27cbf4a1bb04f8fce327a0 Mon Sep 17 00:00:00 2001
From: niksirbi <niko.sirbiladze@gmail.com>
Date: Fri, 18 Oct 2024 15:47:15 +0100
Subject: [PATCH 7/9] fixe mypy errors

---
 brainglobe_template_builder/preproc/alignment.py       | 2 +-
 brainglobe_template_builder/preproc/cropping.py        | 2 +-
 brainglobe_template_builder/preproc/masking.py         | 4 ++--
 brainglobe_template_builder/preproc/transform_utils.py | 2 +-
 pyproject.toml                                         | 3 +++
 5 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/brainglobe_template_builder/preproc/alignment.py b/brainglobe_template_builder/preproc/alignment.py
index 522a836..c031dec 100644
--- a/brainglobe_template_builder/preproc/alignment.py
+++ b/brainglobe_template_builder/preproc/alignment.py
@@ -183,7 +183,7 @@ def _compute_transform(self):
             translation_to_mid_axis @ rotation @ translation_to_origin
         )
 
-    def transform_image(self, image: np.ndarray = None):
+    def transform_image(self, image: np.ndarray | None = None):
         """Transform the image using the transformation matrix.
 
         Parameters
diff --git a/brainglobe_template_builder/preproc/cropping.py b/brainglobe_template_builder/preproc/cropping.py
index f2a4f17..a34cd73 100644
--- a/brainglobe_template_builder/preproc/cropping.py
+++ b/brainglobe_template_builder/preproc/cropping.py
@@ -2,7 +2,7 @@
 
 
 def crop_to_mask(
-    stack: np.ndarray, mask: np.ndarray, padding: np.uint8 = 0
+    stack: np.ndarray, mask: np.ndarray, padding: np.uint8 = np.uint8(0)
 ) -> tuple[np.ndarray, np.ndarray]:
     """
     Crop stack and mask to the mask extent, and pad with zeros.
diff --git a/brainglobe_template_builder/preproc/masking.py b/brainglobe_template_builder/preproc/masking.py
index f81c5e3..6f4582b 100644
--- a/brainglobe_template_builder/preproc/masking.py
+++ b/brainglobe_template_builder/preproc/masking.py
@@ -1,4 +1,4 @@
-from typing import Literal, Union
+from typing import Callable, Literal, Union
 
 import numpy as np
 from skimage import filters, measure, morphology
@@ -44,7 +44,7 @@ def _threshold_image(
         A binary mask.
     """
 
-    method_to_func = {
+    method_to_func: dict[str, Callable] = {
         "triangle": filters.threshold_triangle,
         "otsu": filters.threshold_otsu,
         "isodata": filters.threshold_isodata,
diff --git a/brainglobe_template_builder/preproc/transform_utils.py b/brainglobe_template_builder/preproc/transform_utils.py
index e9bab36..57ee085 100644
--- a/brainglobe_template_builder/preproc/transform_utils.py
+++ b/brainglobe_template_builder/preproc/transform_utils.py
@@ -118,7 +118,7 @@ def downsample_anisotropic_image_stack(
 
     # we have xy slices as chunks, so apply downscaling in xy first
     downsampled_inplane = stack.map_blocks(
-        transform.downscale_local_mean,
+        transform.downscale_local_mean,  # type: ignore
         (1, in_plane_factor, in_plane_factor),
         dtype=np.float64,
     )
diff --git a/pyproject.toml b/pyproject.toml
index f294c67..c10c876 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -83,6 +83,9 @@ ignore = [
     "scripts/",
 ]
 
+[tool.mypy]
+ignore_missing_imports = true
+exclude = 'examples'
 
 [tool.ruff]
 line-length = 79

From c59b7047b4c1af709fbd8feaa3f7c8f5c3523212 Mon Sep 17 00:00:00 2001
From: niksirbi <niko.sirbiladze@gmail.com>
Date: Fri, 18 Oct 2024 16:36:32 +0100
Subject: [PATCH 8/9] force mypy to only run on the actual package

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index c10c876..8e20a2a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -85,7 +85,7 @@ ignore = [
 
 [tool.mypy]
 ignore_missing_imports = true
-exclude = 'examples'
+explicit_package_bases = true
 
 [tool.ruff]
 line-length = 79

From ca15d9fba0d47232c091e223062d24d3a3707d98 Mon Sep 17 00:00:00 2001
From: Alessandro Felder <alessandrofelder@users.noreply.github.com>
Date: Thu, 31 Oct 2024 13:44:46 +0000
Subject: [PATCH 9/9] Update examples/tadpole/2_prep_lowres.py

---
 examples/tadpole/2_prep_lowres.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/tadpole/2_prep_lowres.py b/examples/tadpole/2_prep_lowres.py
index da8707d..13f5339 100644
--- a/examples/tadpole/2_prep_lowres.py
+++ b/examples/tadpole/2_prep_lowres.py
@@ -1,5 +1,5 @@
 """
-Prepare low-resolution BlackCap images for template construction
+Prepare low-resolution tadpole images for template construction
 ================================================================
 The ``brainglobe-template-builder`` Preprocess widget was used to perform
 the following operations on each subject's low-resolution image, ahead of