bug fixes 2

Former-commit-id: e7f4b42
bhklab · Dec 9, 2021 · d523d79 · d523d79
1 parent 2ef1bb6
commit d523d79
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 16 deletions.
diff --git a/imgtools/autopipeline.py b/imgtools/autopipeline.py
@@ -93,8 +93,6 @@ def process_one_subject(self, subject_id):
         print(read_results)
 
         print(subject_id, " start")
-        #For counting multiple connections per modality
-        counter = {"CT":0,"RTDOSE":0,"RTSTRUCT":0,"PT":0}
 
         metadata = {}
         for i, colname in enumerate(self.output_streams):
@@ -104,7 +102,9 @@ def process_one_subject(self, subject_id):
             output_stream = ("_").join([item for item in colname.split("_") if item != "1"])
 
             #If there are multiple connections existing, multiple connections means two modalities connected to one modality. They end with _1
-            mult_conn = colname.split("_")[-1] == "1"
+            mult_conn = colname.split("_")[-1].isnumeric()
+            num = colname.split("_")[-1]
+
             print(output_stream)
 
             if read_results[i] is None:
@@ -136,8 +136,7 @@ def process_one_subject(self, subject_id):
                 if not mult_conn:
                     self.output(subject_id, doses, output_stream)
                 else:
-                    counter[modality] = counter[modality]+1
-                    self.output(f"{subject_id}_{counter[modality]}", doses, output_stream)
+                    self.output(f"{subject_id}_{num}", doses, output_stream)
                 metadata[f"size_{output_stream}"] = str(doses.GetSize())
                 metadata[f"metadata_{colname}"] = [read_results[i].get_metadata()]
                 print(subject_id, " SAVED DOSE")
@@ -158,8 +157,7 @@ def process_one_subject(self, subject_id):
                 if not mult_conn:
                     self.output(subject_id, mask, output_stream)
                 else:
-                    counter[modality] = counter[modality] + 1
-                    self.output(f"{subject_id}_{counter[modality]}", mask, output_stream)
+                    self.output(f"{subject_id}_{num}", mask, output_stream)
                 metadata[f"metadata_{colname}"] = [structure_set.roi_names]
 
                 print(subject_id, "SAVED MASK ON", conn_to)
@@ -174,8 +172,7 @@ def process_one_subject(self, subject_id):
                 if not mult_conn:
                     self.output(subject_id, pet, output_stream)
                 else:
-                    counter[modality] = counter[modality] + 1
-                    self.output(f"{subject_id}_{counter[modality]}", pet, output_stream)
+                    self.output(f"{subject_id}_{num}", pet, output_stream)
                 metadata[f"size_{output_stream}"] = str(pet.GetSize())
                 metadata[f"metadata_{colname}"] = [read_results[i].get_metadata()]
                 print(subject_id, " SAVED PET")

diff --git a/imgtools/io/dataset.py b/imgtools/io/dataset.py
@@ -56,10 +56,10 @@ def load_from_nrrd(
             temp = {}
             for col in output_streams:
                 extension = file_names[col]
-                mult_conn = col.split("_")[-1] == "1"
+                mult_conn = col.split("_")[-1].isnumeric()
                 metadata_name = f"metadata_{col}"
                 if mult_conn:
-                    extra = str(col.split("_").count("1"))+"_"
+                    extra = col.split("_")[-1]+"_"
                 else:
                     extra = ""
                 path_mod = os.path.join(path,extension.split(".")[0],f"{subject_id}_{extra}{extension}.nrrd")
@@ -127,7 +127,7 @@ def process_one_subject(
             output_stream = ("_").join([item for item in colname.split("_") if item != "1"])
 
             if read_results[i] is None:
-                pass
+                temp[f"mod_{colname}"] = None
             elif modality == "CT":
                 image = read_results[i]
                 if len(image.GetSize()) == 4:

diff --git a/imgtools/modules/datagraph.py b/imgtools/modules/datagraph.py
@@ -413,9 +413,11 @@ def _get_df(self,
     def _check_save(save_dict,node,dest):
         key = f"folder_{node}_{dest}"
         key_series = f"series_{node}_{dest}"
-        if key in save_dict.keys():
-            key = key + "_1"
-            key_series = key_series + "_1"
+        i = 1
+        while key in save_dict.keys():
+            key = f"folder_{node}_{dest}_{i}"
+            key_series = f"series_{node}_{dest}_{i}"
+            i+=1
         return key,key_series
 
     @staticmethod

diff --git a/tests/test_autopipe.py b/tests/test_autopipe.py
@@ -91,6 +91,6 @@ def test_pipeline(dataset_path,modalities):
         os.remove(crawl_path)
         os.remove(json_path)
         os.remove(edge_path)
-    shutil.rmtree(output_path_mod)
+    # shutil.rmtree(output_path_mod)