adjusted code

code/register_masks.py

@@ -1,10 +1,3 @@
-# -*- coding: utf-8 -*-
-"""
-Created on Fri Jan 12 18:22:18 2024
-
-@author: arefks
-"""
-
 import argparse
 import subprocess
 import shlex
@@ -13,6 +6,8 @@ import glob
 import csv
 import nibabel as nib
 import numpy as np
+from tqdm import tqdm
+
 def main():
     parser = argparse.ArgumentParser(description="Image Registration and Resampling")
     
@@ -24,17 +19,20 @@ def main():
     missing_matrix_paths = []
 
     # Parsing for all BiasBet.nii.gz files in T2w folders after a given input path.
-    T2BetPath = os.path.join(args.input,"**","T2w","*BiasBet.nii.gz")
+    T2BetPath = os.path.join(args.input,"**","anat","*BiasBet.nii.gz")
     T2BetFiles = glob.glob(T2BetPath,recursive=True)
     
     # Parsing for all given masks.
     FloatingMasksPath = os.path.join(args.masks, "*.nii*")
     FloatingMasksFiles = glob.glob(FloatingMasksPath)
     
+    total_iterations = len(T2BetFiles) * len(FloatingMasksFiles)
+    progress_bar = tqdm(total=total_iterations, desc='Registration Progress (Mask to Anatomical T2w)', unit='files')
+
     for ff in T2BetFiles:
         for mm in FloatingMasksFiles:
             MaskName = os.path.basename(mm).replace(".nii","").replace(".gz","")
-            NewMaskName = MaskName + "_T2w" + ".nii.gz"
+            NewMaskName = MaskName + "_anat" + ".nii.gz"
             NewMaskFolder = os.path.join(os.path.dirname(ff),"RegisteredTractMasks")
             NewMaskAddress = os.path.join(NewMaskFolder,NewMaskName)
             if not os.path.exists(NewMaskFolder):
@@ -56,19 +54,25 @@ def main():
                 print(f"Errors:\n{result.stderr.decode()}")
             except Exception as e:
                 print(f"An error occurred: {e}")
+                
+            progress_bar.update(1)
     
+    progress_bar.close()
     print("Registered all masks to t2w space of individual mice: DONE\n")
     
     # Parsing for all SmoothMicoBet.nii.gz files in DTI folders after a given input path.
-    DTIBetPath = os.path.join(args.input,"**","DTI*","*SmoothMicoBet.nii.gz")
+    DTIBetPath = os.path.join(args.input,"**","dwi*","*SmoothMicoBet.nii.gz")
     DTIBetFiles = glob.glob(DTIBetPath,recursive=True)
     
+    total_iterations = len(DTIBetFiles) * len(FloatingMasksFiles)
+    progress_bar = tqdm(total=total_iterations, desc='Registration Progress (to diffusion dwi) and adjusting for fa, ad, md maps', unit='files')
+
     for ff in DTIBetFiles:
         temp = os.path.dirname(os.path.dirname(ff))
-        t2MasksPath = os.path.join(temp,"T2w","RegisteredTractMasks","*nii.gz")
+        t2MasksPath = os.path.join(temp,"anat","RegisteredTractMasks","*nii.gz")
         FloatingMasksFiles = glob.glob(t2MasksPath)
         for mm in FloatingMasksFiles:
-            MaskName = os.path.basename(mm).replace("T2w","DTI")
+            MaskName = os.path.basename(mm).replace("anat","dwi")
             NewMaskName = MaskName
             NewMaskFolder = os.path.join(os.path.dirname(ff),"RegisteredTractMasks")
             NewMaskAddress = os.path.join(NewMaskFolder,NewMaskName)
@@ -90,12 +94,20 @@ def main():
                 print(f"Errors:\n{result.stderr.decode()}")
             except Exception as e:
                 print(f"An error occurred: {e}")
+                
+            progress_bar.update(1)
     
+    progress_bar.close()
+    print("Registered all masks to t2w space of individual mice: DONE\n")
+
     # Adjust by flipping for the fa0 and ad files
     # Locate all files in the "RegisteredTractMasks" folder
     registered_masks_path = os.path.join(args.input,"**","RegisteredTractMasks", "*.nii.gz")
     registered_masks_files = glob.glob(registered_masks_path, recursive=True)
 
+    total_iterations = len(registered_masks_files)
+    progress_bar = tqdm(total=total_iterations, desc='Flipping Masks', unit='files')
+
     for mask_file in registered_masks_files:
         # Load the mask file
         mask_nifti = nib.load(mask_file)
@@ -118,14 +130,17 @@ def main():
         adjusted_mask_file = os.path.join(adjusted_masks_dir, os.path.basename(mask_file).replace(".nii.gz","_flipped.nii.gz"))
         nib.save(nib.Nifti1Image(flipped_mask, mask_nifti.affine), adjusted_mask_file)
 
-        print("Flipping of RegisteredTractMasks and saving to RegisteredTractMasks_adjusted: DONE\n")
+        progress_bar.update(1)
+
+    progress_bar.close()
+    print("Flipping of RegisteredTractMasks and saving to RegisteredTractMasks_adjusted: DONE\n")
 
     # Save the missing matrix paths to a CSV file
     missing_matrix_csv_path = os.path.join(args.input, "missing_matrix_paths.csv")
     with open(missing_matrix_csv_path, mode='w', newline='') as file:
         writer = csv.writer(file)
         if not missing_matrix_paths:
-            writer.writerow(["All folders contained an Affine transformation matrix"])
+            writer.writerow(["Registration of all masks/maps was executed. All folders contained an Affine transformation matrix"])
         else:
             for path in missing_matrix_paths:
                 writer.writerow([path])
@@ -134,4 +149,4 @@ def main():
 
         
 if __name__ == "__main__":
-    main()
+    main()