split mask up into each contour

imgtools/io/common.py

@@ -43,11 +43,11 @@ def file_name_convention() -> Dict:
     file_name_convention = {"CT": "image",
                             "MR": "image",
                             "RTDOSE_CT": "dose", 
-                            "RTSTRUCT_CT": "mask_ct.seg", 
-                            "RTSTRUCT_MR": "mask_mr.seg", 
-                            "RTSTRUCT_PT": "mask_pt.seg", 
+                            "RTSTRUCT_CT": "mask_ct", 
+                            "RTSTRUCT_MR": "mask_mr", 
+                            "RTSTRUCT_PT": "mask_pt", 
                             "PT_CT": "pet", 
                             "PT": "pet", 
                             "RTDOSE": "dose", 
-                            "RTSTRUCT": "mask.seg"}
+                            "RTSTRUCT": "mask"}
     return file_name_convention

imgtools/io/writers.py

@@ -2,6 +2,7 @@
 import json
 import csv
 import pickle
+import shutil
 from datetime import datetime, timezone
 
 import h5py
@@ -20,8 +21,8 @@ def __init__(self, root_directory, filename_format, create_dirs=True):
         self.filename_format = filename_format
         self.create_dirs = create_dirs
         #this one makes an extra {subject_id} folder
-        # if create_dirs and not os.path.exists(self.root_directory):
-        #     os.makedirs(self.root_directory)
+        if create_dirs and not os.path.exists(self.root_directory):
+            os.makedirs(self.root_directory)
 
     def put(self, *args, **kwargs):
         raise NotImplementedError
@@ -36,8 +37,8 @@ def _get_path_from_subject_id(self, subject_id, **kwargs):
                                                    time=time,
                                                    date_time=date_time,
                                                    **kwargs)
-        self.root_directory = self.root_directory.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, out_filename)
         out_dir = os.path.dirname(out_path)
         if self.create_dirs and not os.path.exists(out_dir):
@@ -46,19 +47,53 @@ 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, **kwargs):        
+        if kwargs["is_mask"]:
+            self.filename_format = kwargs["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, out_filename)
+        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)
         self.compress = compress
 
     def put(self, subject_id, image, **kwargs):
-        # TODO (Michal) add support for .seg.nrrd files
         out_path = self._get_path_from_subject_id(subject_id, **kwargs)
         sitk.WriteImage(image, out_path, self.compress)
 
 
 class SegNrrdWriter(BaseWriter):
     def __init__(self, root_directory, filename_format="{subject_id}.seg.nrrd", create_dirs=True, compress=True):
+        if filename_format.endswith(".nii.gz"):
+            filename_format = filename_format.replace(".nii.gz", ".nrrd") #.seg is already included
         super().__init__(root_directory, filename_format, create_dirs)
         if compress:
             self.compression_level = 9
@@ -67,7 +102,7 @@ def __init__(self, root_directory, filename_format="{subject_id}.seg.nrrd", crea
 
     def put(self, subject_id, mask, **kwargs):
         out_path = self._get_path_from_subject_id(subject_id, **kwargs)
-        labels = [k for k in mask.roi_names] 
+        labels = [k for k in mask.roi_names]
         print(labels)
 
         origin = mask.GetOrigin()

imgtools/modules/structureset.py

@@ -31,6 +31,7 @@ def from_dicom_rtstruct(cls, rtstruct_path: str) -> 'StructureSet':
                 roi_points[name] = _get_roi_points(rtstruct, i)
             except AttributeError:
                 warn(f"Could not get points for ROI {name} (in {rtstruct_path}).")
+        print(roi_names, "\n\n\n")
         return cls(roi_points)
 
     @property
@@ -162,7 +163,7 @@ def to_segmentation(self, reference_image: sitk.Image,
                 #     mask[z, :, :, label] += slice_mask
 
 
-        mask[mask > 1] = 1        
+        mask[mask > 1] = 1
         mask = sitk.GetImageFromArray(mask, isVector=True)
         mask.CopyInformation(reference_image)
         seg_roi_names = {"_".join(k): v for v, k in groupby(labels, key=lambda x: labels[x])}

imgtools/ops/ops.py

@@ -241,7 +241,7 @@ def __init__(self,
         self.create_dirs = create_dirs
         self.compress = compress
 
-        if "seg.nrrd" in filename_format:
+        if ".seg" in filename_format: #from .seg.nrrd bc it is now .nii.gz
             writer_class = SegNrrdWriter
         else:
             writer_class = ImageFileWriter
@@ -253,6 +253,30 @@ def __init__(self,
 
         super().__init__(writer)
 
+
+
+class ImageSubjectFileOutput(BaseOutput):
+
+    def __init__(self,
+                 root_directory: str,
+                 filename_format: Optional[str] ="{subject_id}.nrrd",
+                 create_dirs: Optional[bool] =True,
+                 compress: Optional[bool] =True):
+        self.root_directory = root_directory
+        self.filename_format = filename_format
+        self.create_dirs = create_dirs
+        self.compress = compress
+
+        writer_class = BaseSubjectWriter
+
+        writer = writer_class(self.root_directory,
+                              self.filename_format,
+                              self.create_dirs,
+                              self.compress)
+
+        super().__init__(writer)
+
+
 class ImageAutoOutput:
     """
     Wrapper class around ImageFileOutput. This class supports multiple modalities writers and calls ImageFileOutput for writing the files
@@ -280,17 +304,25 @@ def __init__(self,
             # Not considering colnames ending with alphanumeric
             colname_process = ("_").join([item for item in colname.split("_") if item.isnumeric()==False])
             extension = self.file_name[colname_process]
-            self.output[colname_process] = ImageFileOutput(os.path.join(root_directory,"{subject_id}",extension.split(".")[0]),
-                                                           filename_format=colname_process+"{}.nii.gz".format(extension))
+            self.output[colname_process] = ImageSubjectFileOutput(os.path.join(root_directory,"{subject_id}",extension.split(".")[0]),
+                                                                    filename_format=colname_process+"{}.nii.gz".format(extension))
             # self.output[colname_process] = ImageFileOutput(os.path.join(root_directory,extension.split(".")[0]),
             #                                                filename_format="{subject_id}_"+"{}.nrrd".format(extension))
+            # if not is_mask:
+            #         self.output[colname_process] = ImageSubjectFileOutput(os.path.join(root_directory,"{subject_id}",extension.split(".")[0]),
+            #                                                               filename_format=colname_process+"{}.nii.gz".format(extension))
+            # else:
+            #     self.output[colname_process] = ImageSubjectFileOutput(os.path.join(root_directory,"{subject_id}",extension.split(".")[0]),
+            #                                                             filename_format=mask_label+"{}.nii.gz".format(extension),)
 
     def __call__(self, 
                  subject_id: str,
                  img: sitk.Image,
-                 output_stream):
+                 output_stream,
+                 is_mask: bool = False,
+                 mask_label: Optional[str] = ""):
 
-        self.output[output_stream](subject_id, img)
+        self.output[output_stream](subject_id, img, is_mask=is_mask, mask_label=mask_label)
 
 class NumpyOutput(BaseOutput):
     """NumpyOutput class processed images as NumPy files.