Merge pull request #41 from Hyedryn/tracking_mask

Tracking mask
Hyedryn · Apr 11, 2024 · 0be8e87 · 0be8e87
2 parents 26db8ab + 1932e49
commit 0be8e87
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 22 deletions.
diff --git a/elikopy/core.py b/elikopy/core.py
@@ -989,7 +989,7 @@ def odf_msmtcsd(self, folder_path=None, num_peaks = 2, peaks_threshold=0.25, mas
         f.close()
 
 
-    def tracking(self, folder_path=None, streamline_number:int=100000, max_angle:int=15, cutoff:float=0.1, msmtCSD:bool=True, output_filename:str="tractogram", save_as_trk=False,
+    def tracking(self, folder_path=None, streamline_number:int=100000, max_angle:int=15, cutoff:float=0.1, msmtCSD:bool=True, output_filename:str="tractogram", save_as_trk=False, maskType:str="brain_mask",
                  slurm=None, patient_list_m=None, slurm_email=None, slurm_timeout=None, cpus=None, slurm_mem=None):
         """Computes the odf using MSMT-CSD for each subject. The outputs are available in the directories <folder_path>/subjects/<subjects_ID>/dMRI/tractography/.
 
@@ -1034,7 +1034,7 @@ def tracking(self, folder_path=None, streamline_number:int=100000, max_angle:int
 
             if slurm:
                 p_job = {
-                        "wrap": "export MKL_NUM_THREADS="+ str(core_count)+" ; export OMP_NUM_THREADS="+ str(core_count)+" ; python -c 'from elikopy.individual_subject_processing import tracking_solo; tracking_solo(\"" + folder_path + "/\",\"" + p + "\", streamline_number =" + str(streamline_number) + ", msmtCSD=" + str(msmtCSD) + ", core_count=" + str(core_count) + ", save_as_trk=" + str(save_as_trk) + ", cutoff=" + str(cutoff) + ", max_angle="+ str(max_angle) + ", output_filename= \"" + output_filename + "\"" + ")'",
+                        "wrap": "export MKL_NUM_THREADS="+ str(core_count)+" ; export OMP_NUM_THREADS="+ str(core_count)+" ; python -c 'from elikopy.individual_subject_processing import tracking_solo; tracking_solo(\"" + folder_path + "/\",\"" + p + "\", streamline_number =" + str(streamline_number) + ", msmtCSD=" + str(msmtCSD) + ", core_count=" + str(core_count) + ", save_as_trk=" + str(save_as_trk) + ", cutoff=" + str(cutoff) + ", max_angle="+ str(max_angle) + ", maskType= \"" + maskType + "\"" + ", output_filename= \"" + output_filename + "\"" + ")'",
                         "job_name": "TRACKING_" + p,
                         "ntasks": 1,
                         "cpus_per_task": core_count,
@@ -1054,7 +1054,7 @@ def tracking(self, folder_path=None, streamline_number:int=100000, max_angle:int
                 f.write("["+log_prefix+"] " + datetime.datetime.now().strftime("%d.%b %Y %H:%M:%S") + ": Patient %s is ready to be processed\n" % p)
                 f.write("["+log_prefix+"] " + datetime.datetime.now().strftime("%d.%b %Y %H:%M:%S") + ": Successfully submited job %s using slurm\n" % p_job_id)
             else:
-                tracking_solo(folder_path + "/", p,  streamline_number=streamline_number, max_angle=max_angle, cutoff=cutoff, msmtCSD=msmtCSD, output_filename=output_filename, core_count=core_count, save_as_trk=save_as_trk)
+                tracking_solo(folder_path + "/", p,  streamline_number=streamline_number, max_angle=max_angle, cutoff=cutoff, msmtCSD=msmtCSD, output_filename=output_filename, core_count=core_count, save_as_trk=save_as_trk, maskType=maskType)
                 matplotlib.pyplot.close(fig='all')
                 f.write("["+log_prefix+"] " + datetime.datetime.now().strftime("%d.%b %Y %H:%M:%S") + ": Successfully applied tracking_solo on patient %s\n" % p)
                 f.flush()

diff --git a/elikopy/individual_subject_processing.py b/elikopy/individual_subject_processing.py
@@ -3591,7 +3591,8 @@ def print_page(self, images):
 
 def tracking_solo(folder_path:str, p:str, streamline_number:int=100000,
                   max_angle:int=15, cutoff:float=0.1, msmtCSD:bool=True,
-                  output_filename:str='tractogram',core_count:int=1, save_as_trk=False):
+                  output_filename:str='tractogram',core_count:int=1,
+                  maskType:str="brain_mask",save_as_trk=False):
     """ Computes the whole brain tractogram of a single patient based on the fod obtained from msmt-CSD.
 
     :param folder_path: the path to the root directory.
@@ -3601,16 +3602,19 @@ def tracking_solo(folder_path:str, p:str, streamline_number:int=100000,
     :param cutoff: Value below which streamlines do not propagate. default=0.1
     :param msmtCSD: boolean. If True then uses ODF from msmt-CSD, if False from CSD. default=True
     :param output_filename: str. Specify output filename for tractogram.
+    :param maskType: str. Specify a masking region of interest, streamlines exiting the mask will be truncated.
     """
-
+
+    assert maskType in ["brain_mask_dilated","brain_mask"], "The mask parameter must be one of the following : brain_mask_dilated, brain_mask"
+
     import nibabel as nib
     from elikopy.utils import dipy_fod_to_mrtrix
     from dipy.io.streamline import load_tractogram, save_trk
-    
+
     patient_path = p
-    
-    params={'Number of streamlines':streamline_number, 'Maximum angle':max_angle,
-            'Cutoff value':cutoff}
+
+    params={'Number of streamlines': streamline_number, 'Maximum angle': max_angle,
+            'Cutoff value': cutoff, 'Mask type': maskType}
 
     if msmtCSD:
         if not os.path.isdir(folder_path + '/subjects/' + patient_path + "/dMRI/ODF/MSMT-CSD/"):
@@ -3628,38 +3632,40 @@ def tracking_solo(folder_path:str, p:str, streamline_number:int=100000,
         odf_file_path = folder_path + '/subjects/' + patient_path + "/dMRI/ODF/CSD/"+patient_path + "_CSD_SH_ODF_mrtrix.nii.gz"
         params['Local modeling']='CSD'
     tracking_path = folder_path + '/subjects/' + patient_path + "/dMRI/tractography/"
-    mask_path = folder_path + '/subjects/' + patient_path + '/masks/' + patient_path + "_brain_mask_dilated.nii.gz"
+    seed_path = folder_path + '/subjects/' + patient_path + '/masks/' + patient_path + "_brain_mask.nii.gz"
+    mask_path = folder_path + '/subjects/' + patient_path + '/masks/' + patient_path + '_' + maskType + '.nii.gz'
     dwi_path = folder_path + '/subjects/' + patient_path + '/dMRI/preproc/' + patient_path + '_dmri_preproc.nii.gz'
 
     output_file = tracking_path+patient_path+'_'+output_filename+'.tck'
-    
+
     if not os.path.isdir(tracking_path):
         os.mkdir(tracking_path)
-    
+
     bashCommand=('tckgen -nthreads ' + str(core_count) + ' ' + odf_file_path +' '+ output_file+
-                 ' -seed_image ' +mask_path+
+                 ' -seed_image ' +seed_path+
                  ' -select ' +str(streamline_number)+
                  ' -angle ' +str(max_angle)+
                  ' -cutoff ' +str(cutoff)+
+                 ' -mask ' +mask_path+
                  ' -force')
-    
+
     tracking_log = open(tracking_path+"tractography_logs.txt", "a+")
-    process = subprocess.Popen(bashCommand, universal_newlines=True, shell=True, stdout=tracking_log,
-                               stderr=subprocess.STDOUT)
+    process = subprocess.Popen(bashCommand, universal_newlines=True, shell=True,
+                               stdout=tracking_log, stderr=subprocess.STDOUT)
 
     process.communicate()
 
     tracking_log.close()
-    
+
     if save_as_trk:
         tract = load_tractogram(output_file, dwi_path)
-    
+
         save_trk(tract, output_file[:-3]+'trk')
-    
+
     with open(output_file[:-3]+'json', 'w') as outfile:
         json.dump(params, outfile)
-        
-        
+
+
 def sift_solo(folder_path: str, p: str, streamline_number: int = 100000,
               msmtCSD: bool = True, input_filename: str = 'tractogram',
               core_count: int = 1, save_as_trk=False):

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "elikopy"
-version = "0.4.3"
+version = "0.4.4"
 description = "A set of tools for analysing dMRI"
 authors = [
     "qdessain <quentin.dessain@uclouvain.be>",