added sparse mask class and generating function

Former-commit-id: 8dab6bc

examples/autotest.py

@@ -58,7 +58,7 @@ def __init__(self,
 
         # image processing ops
         self.resample = Resample(spacing=self.spacing)
-        self.make_binary_mask = StructureSetToSegmentation(roi_names=[], continuous=False)
+        self.make_binary_mask = StructureSetToSegmentation(roi_names=[], continuous=False) # "GTV-.*"
 
         # output ops
         self.output = ImageAutoOutput(self.output_directory, self.output_streams)
@@ -154,6 +154,11 @@ def process_one_subject(self, subject_id):
                 # save output
                 print(mask.GetSize())
                 mask_arr = np.transpose(sitk.GetArrayFromImage(mask))
+
+                # if there is only one ROI, sitk.GetArrayFromImage() will return a 3d array instead of a 4d array with one slice
+                if len(mask_arr.shape) == 3:
+                    mask_arr = mask_arr.reshape(1, mask_arr.shape[0], mask_arr.shape[1], mask_arr.shape[2])
+
                 print(mask_arr.shape)
                 roi_names_list = list(mask.roi_names.keys())
                 for i in range(mask_arr.shape[0]):
@@ -219,11 +224,16 @@ def run(self):
 
 
 if __name__ == "__main__":
-    pipeline = AutoPipeline(input_directory="C:/Users/qukev/BHKLAB/hnscc_testing/HNSCC",
-                            output_directory="C:/Users/qukev/BHKLAB/hnscc_testing_output",
+    pipeline = AutoPipeline(input_directory="C:/Users/qukev/BHKLAB/datasetshort/manifest-1598890146597/NSCLC-Radiomics-Interobserver1",
+                            output_directory="C:/Users/qukev/BHKLAB/autopipelineoutputshort",
                             modalities="CT,RTSTRUCT",
                             visualize=False)
 
+    # pipeline = AutoPipeline(input_directory="C:/Users/qukev/BHKLAB/hnscc_testing/HNSCC",
+    #                         output_directory="C:/Users/qukev/BHKLAB/hnscc_testing_output",
+    #                         modalities="CT,RTSTRUCT",
+    #                         visualize=False)
+
     print(f'starting Pipeline...')
     pipeline.run()
 

imgtools/modules/sparsemask.py

@@ -0,0 +1,8 @@
+from typing import Dict
+
+import numpy as np
+
+class SparseMask:
+    def __init__(self, mask_array:np.ndarray, roi_name_dict: Dict[str, int]):
+        self.mask_array = mask_array
+        self.roi_name_dict = roi_name_dict

imgtools/modules/structureset.py

@@ -1,6 +1,6 @@
 import re
 from warnings import warn
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Tuple, Set
 
 import numpy as np
 import SimpleITK as sitk
@@ -9,6 +9,7 @@
 from skimage.draw import polygon2mask
 
 from .segmentation import Segmentation
+from .sparsemask import SparseMask
 from ..utils import physical_points_to_idxs
 
 
@@ -50,7 +51,7 @@ def _assign_labels(self, names, force_missing=False):
             for i, name in enumerate(self.roi_names):
                 labels[name] = i
         else:
-            for j, pat in enumerate(names):
+            for _, pat in enumerate(names):
                 if sorted(names) == sorted(list(labels.keys())): #checks if all ROIs have already been processed.
                     break
                 if isinstance(pat, str):
@@ -171,6 +172,66 @@ def to_segmentation(self, reference_image: sitk.Image,
         mask = Segmentation(mask, roi_names=seg_roi_names)
 
         return mask
+
+    def generate_sparse_mask(self, mask: Segmentation) -> SparseMask:
+        """
+        Generate a sparse mask from the contours, taking the argmax of all overlaps
+
+        Parameters
+        ----------
+        mask
+            Segmentation object to build sparse mask from
+
+        Returns
+        -------
+        SparseMask
+            The sparse mask object.
+        """
+        mask_arr = np.transpose(sitk.GetArrayFromImage(mask))
+        roi_names = {k: v+1 for k, v in mask.roi_names.items()}
+
+        sparsemask_arr = np.zeros(mask_arr.shape[1:])
+
+        # voxels_with_overlap = {}
+        for i in len(mask_arr.shape[0]):
+            slice = mask_arr[i, :, :, :]
+            slice *= list(roi_names.values())[i] # everything is 0 or 1, so this is fine to convert filled voxels to label indices
+            # res = self._max_adder(sparsemask_arr, slice)
+            # sparsemask_arr = res[0]
+            # for e in res[1]:
+            #     voxels_with_overlap.add(e)
+            sparsemask_arr = np.fmax(sparsemask_arr, slice) # elementwise maximum
+
+        sparsemask = SparseMask(sparsemask_arr, roi_names)
+        # if len(voxels_with_overlap) != 0:
+        #     raise Warning(f"{len(voxels_with_overlap)} voxels have overlapping contours.")
+        return sparsemask
+
+    def _max_adder(self, arr_1: np.ndarray, arr_2: np.ndarray) -> Tuple[np.ndarray, Set[Tuple[int, int, int]]]:
+        """
+        Takes the maximum of two 3D arrays elementwise and returns the resulting array and a list of voxels that have overlapping contours in a set
+
+        Parameters
+        ----------
+        arr_1
+            First array to take maximum of
+        arr_2
+            Second array to take maximum of
+        
+        Returns
+        -------
+        Tuple[np.ndarray, Set[Tuple[int, int, int]]]
+            The resulting array and a list of voxels that have overlapping contours in a set
+        """
+        res = np.zeros(arr_1.shape)
+        overlaps = {} #set of tuples of the coords that have overlap
+        for i in range(arr_1.shape[0]):
+            for j in range(arr_1.shape[1]):
+                for k in range(arr_1.shape[2]):
+                    if arr_1[i, j, k] != 0 and arr_2[i, j, k] != 0:
+                        overlaps.add((i, j, k))
+                    res[i, j, k] = max(arr_1[i, j, k], arr_2[i, j, k])
+        return res, overlaps
 
     def __repr__(self):
         return f"<StructureSet with ROIs: {self.roi_names!r}>"