brought back basesubjectwriter

Former-commit-id: 9a8ff65
bhklab · May 19, 2022 · eb39ab0 · eb39ab0
1 parent 3a337ac
commit eb39ab0
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/imgtools/io/writers.py b/imgtools/io/writers.py
@@ -44,6 +44,38 @@ def _get_path_from_subject_id(self, subject_id, **kwargs):
         return out_path
 
 
+class BaseSubjectWriter(BaseWriter):
+    def __init__(self, root_directory, filename_format="{subject_id}.nii.gz", create_dirs=True, compress=True):
+        super().__init__(root_directory, filename_format, create_dirs)
+        self.root_directory = root_directory
+        self.filename_format = filename_format
+        self.create_dirs = create_dirs
+        self.compress = compress
+        if os.path.exists(self.root_directory)\
+           and os.path.basename(os.path.dirname(self.root_directory)) == "{subject_id}":
+           #delete the folder called {subject_id} that was made in the original BaseWriter
+
+            shutil.rmtree(os.path.dirname(self.root_directory))
+        print(self.root_directory)
+
+    def put(self, subject_id, image, is_mask=False, mask_label="",**kwargs):
+        if is_mask:
+            self.filename_format = mask_label+".nii.gz"
+        out_path = self._get_path_from_subject_id(subject_id, **kwargs)
+        sitk.WriteImage(image, out_path, self.compress)
+
+    def _get_path_from_subject_id(self, subject_id, **kwargs):
+        # out_filename = self.filename_format.format(subject_id=subject_id, **kwargs)
+        self.root_directory = self.root_directory.format(subject_id=subject_id,
+                                                         **kwargs)
+        out_path = os.path.join(self.root_directory, self.filename_format)
+        out_dir = os.path.dirname(out_path)
+        if self.create_dirs and not os.path.exists(out_dir):
+            os.makedirs(out_dir, exist_ok=True) # create subdirectories if specified in filename_format
+
+        return out_path
+
+
 class ImageFileWriter(BaseWriter):
     def __init__(self, root_directory, filename_format="{subject_id}.nii.gz", create_dirs=True, compress=True):
         super().__init__(root_directory, filename_format, create_dirs)

diff --git a/imgtools/utils/crawl.py b/imgtools/utils/crawl.py
@@ -115,14 +115,19 @@ def crawl(top,
 
     # save one level above imaging folders
     parent, dataset  = os.path.split(top)
+
+    parent_imgtools = os.path.join(parent, "imgtools")
+
+    if not os.path.exists(parent_imgtools):
+        os.makedirs(parent_imgtools)
 
     # save as json
-    with open(os.path.join(parent, f'imgtools_{dataset}.json'), 'w') as f:
+    with open(os.path.join(parent_imgtools, f'imgtools_{dataset}.json'), 'w') as f:
         json.dump(database_dict, f, indent=4)
 
     # save as dataframe
     df = to_df(database_dict)
-    df_path = os.path.join(parent, f'imgtools_{dataset}.csv')
+    df_path = os.path.join(parent_imgtools, f'imgtools_{dataset}.csv')
     df.to_csv(df_path)
 
     return database_dict