make sure all files for template building contain subject ID

brainglobe_template_builder/preproc/splitting.py

@@ -32,12 +32,15 @@ def get_right_and_left_slices(array: np.ndarray) -> tuple:
 
 
 def generate_arrays_4template(
-    brain: np.ndarray, mask: np.ndarray, pad: int = 0
+    subject: str, brain: np.ndarray, mask: np.ndarray, pad: int = 0
 ) -> dict[str, np.ndarray]:
     """Generate all needed arrays for the template building process.
 
     Parameters
     ----------
+    subject : str
+        Subject ID to use as a prefix for the keys in the output dictionary,
+        e.g. "sub-01".
     brain : np.ndarray
         The aligned brain image to split into hemispheres and symmetrise.
     mask : np.ndarray
@@ -51,7 +54,9 @@ def generate_arrays_4template(
 
     Returns
     -------
-    A dictionary containing the following arrays:
+    A dictionary containing the multiple arrays needed for the template
+    building process. The keys start with the subject ID and end with the
+    following suffixes (e.g. "sub-01_asym-brain"):
     - asym-brain: the input aligned image (asymmetric brain)
     - asym-mask: the input aligned mask (asymmetric mask)
     - right-hemi-brain: the right hemisphere of the image
@@ -77,8 +82,8 @@ def generate_arrays_4template(
 
     # Put the input arrays into the dictionary
     out_dict = {
-        "asym-brain": brain,
-        "asym-mask": mask,
+        f"{subject}_asym-brain": brain,
+        f"{subject}_asym-mask": mask,
     }
 
     right_half, left_half = get_right_and_left_slices(brain)
@@ -92,12 +97,12 @@ def generate_arrays_4template(
         left_sym_arr = np.dstack([left_half_arr_xflip, left_half_arr])
         out_dict.update(
             {
-                f"right-hemi-{label}": right_half_arr,
-                f"right-hemi-xflip-{label}": right_half_arr_xflip,
-                f"left-hemi-{label}": left_half_arr,
-                f"left-hemi-xflip-{label}": left_half_arr_xflip,
-                f"right-sym-{label}": right_sym_arr,
-                f"left-sym-{label}": left_sym_arr,
+                f"{subject}_right-hemi-{label}": right_half_arr,
+                f"{subject}_right-hemi-xflip-{label}": right_half_arr_xflip,
+                f"{subject}_left-hemi-{label}": left_half_arr,
+                f"{subject}_left-hemi-xflip-{label}": left_half_arr_xflip,
+                f"{subject}_right-sym-{label}": right_sym_arr,
+                f"{subject}_left-sym-{label}": left_sym_arr,
             }
         )
 

examples/BlackCap_prep_lowres.py

@@ -112,10 +112,10 @@
 for idx, row in tqdm(df.iterrows(), total=n_subjects):
     # Figure out input-output paths
     row = df.iloc[idx]
-    subject_str = "sub-" + row["subject_id"]
+    subject = "sub-" + row["subject_id"]
     channel_str = "channel-" + row["color"]
-    file_prefix = f"{subject_str}_{res_str}_{channel_str}"
-    deriv_subj_dir = deriv_dir / subject_str
+    file_prefix = f"{subject}_{res_str}_{channel_str}"
+    deriv_subj_dir = deriv_dir / subject
     tiff_path = deriv_subj_dir / f"{file_prefix}.tif"
 
     logger.info(f"Starting to process {file_prefix}...")
@@ -234,13 +234,18 @@
 
     # 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(
-        aligned_image.numpy(), aligned_mask.numpy(), pad=2
+        subject, aligned_image.numpy(), aligned_mask.numpy(), pad=2
     )
     save_array_dict_to_nii(array_dict, use4template_dir, vox_sizes)
-    use4template_dirs[subject_str] = use4template_dir
+    use4template_dirs[subject] = use4template_dir
     logger.info(
         f"Saved images for template construction in {use4template_dir}."
     )
@@ -262,7 +267,7 @@
 #    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 = {
+filepath_lists: dict[str, list] = {
     "asym-brain": [],
     "asym-mask": [],
     "sym-brain": [],
@@ -272,36 +277,36 @@
 }
 
 for _, row in df.iterrows():
-    subject_str = "sub-" + row["subject_id"]
-    use4template_dir = use4template_dirs[subject_str]
+    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"asym-{label}.nii.gz"
+                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"right-sym-{label}.nii.gz"
+                use4template_dir / f"{subject}_right-sym-{label}.nii.gz"
             )
             # Add paths for the hemispheric template
             filepath_lists[f"hemi-{label}"].append(
-                use4template_dir / f"right-hemi-{label}.nii.gz"
+                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"left-sym-{label}.nii.gz"
+                use4template_dir / f"{subject}_left-sym-{label}.nii.gz"
             )
             # Add paths for the hemispheric template
             filepath_lists[f"hemi-{label}"].append(
-                use4template_dir / f"left-hemi-xflip-{label}.nii.gz"
+                use4template_dir / f"{subject}_left-hemi-xflip-{label}.nii.gz"
             )
 
 # %%