DICOM-SEG beta commit (#74)

imgtools/io/loaders.py

@@ -7,28 +7,22 @@
 
 import pandas as pd
 import SimpleITK as sitk
-import nrrd
 from pydicom import dcmread
 
 from joblib import Parallel, delayed
 from tqdm.auto import tqdm
 
-from ..modules import StructureSet
-from ..modules import Dose
-from ..modules import PET
-from ..modules import Scan
+from ..modules import StructureSet, Dose, PET, Scan, Segmentation
 from ..utils.crawl import *
 from ..utils.dicomutils import *
 
 def read_image(path):
     return sitk.ReadImage(path)
 
-def read_header(path):
-    return nrrd.read_header(path)
-
 def read_dicom_series(path: str,
                       series_id: Optional[str] = None,
-                      recursive: bool = False) -> Scan:
+                      recursive: bool = False, 
+                      file_names: list = None):
     """Read DICOM series as SimpleITK Image.
 
     Parameters
@@ -45,18 +39,23 @@ def read_dicom_series(path: str,
        the directory. If None and multiple series are present, loads the first
        series found.
 
+    file_names, optional
+        If there are multiple acquisitions/"subseries" for an individual series,
+        use the provided list of file_names to set the ImageSeriesReader.
 
     Returns
     -------
     The loaded image.
 
     """
     reader = sitk.ImageSeriesReader()
-    dicom_names = reader.GetGDCMSeriesFileNames(path,
-                                                seriesID=series_id if series_id else "",
-                                                recursive=recursive)
-    # extract the names of the dicom files that are in the path variable, which is a directory
-    reader.SetFileNames(dicom_names)
+    if file_names is None:
+        file_names = reader.GetGDCMSeriesFileNames(path,
+                                                    seriesID=series_id if series_id else "",
+                                                    recursive=recursive)
+        # extract the names of the dicom files that are in the path variable, which is a directory
+
+    reader.SetFileNames(file_names)
 
     # Configure the reader to load all of the DICOM tags (public+private):
     # By default tags are not loaded (saves time).
@@ -66,19 +65,26 @@ def read_dicom_series(path: str,
     reader.MetaDataDictionaryArrayUpdateOn()
     reader.LoadPrivateTagsOn()
 
-    metadata = {}
+    return reader.Execute()
 
-    return Scan(reader.Execute(), metadata)
+
 
+def read_dicom_scan(path, series_id=None, recursive: bool=False) -> Scan:
+    image = read_dicom_series(path, series_id=series_id, recursive=recursive)
+    return Scan(image, {})
 
 def read_dicom_rtstruct(path):
     return StructureSet.from_dicom_rtstruct(path)
 
 def read_dicom_rtdose(path):
     return Dose.from_dicom_rtdose(path)
 
-def read_dicom_pet(path,series=None):
-    return PET.from_dicom_pet(path,series, "SUV")
+def read_dicom_pet(path, series=None):
+    return PET.from_dicom_pet(path, series, "SUV")
+
+def read_dicom_seg(path, meta, series=None):
+    seg_img = read_dicom_series(path, series)
+    return Segmentation.from_dicom_seg(seg_img, meta)
 
 def read_dicom_auto(path, series=None):
     if path is None:
@@ -88,41 +94,28 @@ def read_dicom_auto(path, series=None):
         meta = dcmread(dcm)
         if meta.SeriesInstanceUID != series and series is not None:
             continue
+
         modality = meta.Modality
-        all_modality_metadata = all_modalities_metadata(meta)
-        if modality == 'CT' or modality == 'MR':
-            dicom_series = read_dicom_series(path,series)#, modality=modality)
-            if modality == 'CT':
-                dicom_series.metadata.update(ct_metadata(meta))
-                dicom_series.metadata.update(all_modality_metadata)
-            else:
-                dicom_series.metadata.update(mr_metadata(meta))
-                dicom_series.metadata.update(all_modality_metadata)
-            return dicom_series
+        if modality in ['CT', 'MR']:
+            obj = read_dicom_scan(path, series)
         elif modality == 'PT':
-            pet = read_dicom_pet(path,series)
-            pet.metadata.update(pet_metadata(meta))
-            pet.metadata.update(all_modality_metadata)
-            return pet
+            obj = read_dicom_pet(path, series)
         elif modality == 'RTSTRUCT':
-            rtstruct = read_dicom_rtstruct(dcm)
-            rtstruct.metadata.update(rtstruct_metadata(meta))
-            rtstruct.metadata.update(all_modality_metadata)
-            return rtstruct
+            obj = read_dicom_rtstruct(dcm)
         elif modality == 'RTDOSE':
-            rtdose = read_dicom_rtdose(dcm)
-            rtdose.metadata.update(all_modality_metadata)
-            return rtdose
+            obj = read_dicom_rtdose(dcm)
+        elif modality == 'SEG':
+            obj = read_dicom_seg(path, meta, series)
         else:
-            if len(dcms)==1:
+            if len(dcms) == 1:
+                print(modality, 'at', dcms[0], 'is NOT implemented yet.')
                 raise NotImplementedError
             else:
                 print("There were no dicoms in this path.")
                 return None
-
-def read_segmentation(path):
-    # TODO read seg.nrrd
-    pass
+
+        obj.metadata.update(get_modality_metadata(meta, modality))
+        return obj
 
 class BaseLoader:
     def __getitem__(self, subject_id):
@@ -149,16 +142,12 @@ def get(self, subject_id, default=None):
 class ImageCSVLoader(BaseLoader):
     def __init__(self,
                  csv_path_or_dataframe,
-                 colnames=None,
-                 seriesnames=None,
+                 colnames=[],
+                 seriesnames=[],
                  id_column=None,
                  expand_paths=False,
                  readers=None):
 
-        if colnames is None:
-            colnames = []
-        if seriesnames is None:
-            seriesnames = []
         if readers is None:
             readers = [read_image] # no mutable defaults https://florimond.dev/en/posts/2018/08/python-mutable-defaults-are-the-source-of-all-evil/
 
@@ -176,7 +165,7 @@ def __init__(self,
             self.paths = csv_path_or_dataframe
             if id_column:
                 self.paths = self.paths.set_index(id_column)
-            if not self.colnames:
+            if len(self.colnames) == 0:
                 self.colnames = self.paths.columns
         else:
             raise ValueError(f"Expected a path to csv file or pd.DataFrame, not {type(csv_path_or_dataframe)}.")

imgtools/modules/datagraph.py

@@ -53,7 +53,11 @@ def form_graph(self):
         '''
         Forms edge table based on the crawled data
         '''
-        #Get reference_rs information from RTDOSE-RTPLAN connections
+        # enforce string type to all columns to prevent dtype merge errors for empty columns
+        for col in self.df:
+            self.df[col] = self.df[col].astype(str)
+
+        #Get reference_rs information from RTDOSE-RTPLAN connections    
         df_filter = pd.merge(self.df, self.df[["instance_uid","reference_rs"]], 
                              left_on="reference_pl", 
                              right_on="instance_uid", 
@@ -191,7 +195,7 @@ def _form_edge_study(self, df, all_study, study_id):
         ct = df_study.loc[df_study["modality"] == "CT"]
         mr = df_study.loc[df_study["modality"] == "MR"]
         pet = df_study.loc[df_study["modality"] == "PT"]
-
+        seg = df_study.loc[df_study["modality"] == "SEG"]
         edge_types = np.arange(7)
         for edge in edge_types:
             if edge==0:    # FORMS RTDOSE->RTSTRUCT, can be formed on both series and instance uid
@@ -217,7 +221,12 @@ def _form_edge_study(self, df, all_study, study_id):
 
             elif edge==5: 
                 df_combined = pd.merge(plan, dose, left_on="instance", right_on="reference_pl")
-
+
+            elif edge==7:  # FORMS RTSTRUCT->CT on ref_ct to series
+                df_ct_seg = pd.merge(ct, seg, left_on="series", right_on="reference_ct")
+                df_mr_seg = pd.merge(mr, seg, left_on="series", right_on="reference_ct")
+                df_combined = pd.concat([df_ct_seg, df_mr_seg])
+
             else:
                 df_combined = pd.merge(struct, plan, left_on="instance", right_on="reference_rs")
 
@@ -257,7 +266,7 @@ def parser(self, query_string: str) -> pd.DataFrame:
         '''
         #Basic processing of just one modality
         supp_mods   = ["RTDOSE", "RTSTRUCT", "CT", "PT", 'MR']
-        edge_def    = {"RTSTRUCT,RTDOSE" : 0, "CT,RTDOSE" : 1, "CT,RTSTRUCT" : 2, "PET,RTSTRUCT" : 3, "CT,PT" : 4, 'MR,RTSTRUCT': 2, "RTPLAN,RTSTRUCT": 6, "RTPLAN,RTDOSE": 5}
+        edge_def    = {"RTSTRUCT,RTDOSE" : 0, "CT,RTDOSE" : 1, "CT,RTSTRUCT" : 2, "PET,RTSTRUCT" : 3, "CT,PT" : 4, 'MR,RTSTRUCT': 2, "RTPLAN,RTSTRUCT": 6, "RTPLAN,RTDOSE": 5, "CT,SEG": 7}
         self.mods   = query_string.split(",")
         self.mods_n = len(self.mods)
 

imgtools/modules/segmentation.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import SimpleITK as sitk
+from pydicom import dcmread
 
 from .sparsemask import SparseMask
 
@@ -35,7 +36,13 @@ def map_over_labels(segmentation, f, include_background=False, return_segmentati
 
 
 class Segmentation(sitk.Image):
-    def __init__(self, segmentation, roi_indices=None, existing_roi_indices=None, raw_roi_names=None):
+    def __init__(self, 
+                 segmentation, 
+                 metadata: Optional[dict] = {}, 
+                 roi_indices=None, 
+                 existing_roi_indices=None, 
+                 raw_roi_names={},
+                 frame_groups=None):
         """Initializes the Segmentation class
         
         Parameters
@@ -48,23 +55,28 @@ def __init__(self, segmentation, roi_indices=None, existing_roi_indices=None, ra
 
         raw_roi_names
             Dictionary of {"ROI": original countor names}
+
+        frame_groups
+            PerFrameFunctionalGroupsSequence (5200, 9230) DICOM metadata
         """
         super().__init__(segmentation)
         self.num_labels = self.GetNumberOfComponentsPerPixel()
+        self.raw_roi_names = raw_roi_names
+        self.metadata = metadata
+        self.frame_groups = frame_groups
+
         if not roi_indices:
             self.roi_indices = {f"label_{i}": i for i in range(1, self.num_labels+1)}
         else:
             self.roi_indices = roi_indices
             if 0 in self.roi_indices.values():
                 self.roi_indices = {k : v+1 for k, v in self.roi_indices.items()}
-        if not raw_roi_names:
-            raw_roi_names={}
-        else:
-            self.raw_roi_names = raw_roi_names
+
         if len(self.roi_indices) != self.num_labels:
             for i in range(1, self.num_labels+1):
                 if i not in self.roi_indices.values():
                     self.roi_indices[f"label_{i}"] = i
+
         self.existing_roi_indices = existing_roi_indices
 
     def get_label(self, label=None, name=None, relabel=False):
@@ -174,4 +186,25 @@ def _max_adder(self, arr_1: np.ndarray, arr_2: np.ndarray) -> Tuple[np.ndarray,
                     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
+        return res, overlaps
+
+    @classmethod
+    def from_dicom_seg(cls, mask, meta):
+        #get duplicates
+        label_counters = {i.SegmentLabel: 1 for i in meta.SegmentSequence}
+        raw_roi_names  = {} #{i.SegmentLabel: i.SegmentNumber for n, i in meta.SegmentSequence}
+        for n, i in enumerate(meta.SegmentSequence):
+            label = i.SegmentLabel
+            num   = i.SegmentNumber
+
+            if label not in raw_roi_names:
+                raw_roi_names[label] = num
+            else:
+                raw_roi_names[f"{label}_{label_counters[label]}"] = num
+                label_counters[label] += 1
+
+
+        frame_groups  = meta.PerFrameFunctionalGroupsSequence
+        return cls(mask, raw_roi_names=raw_roi_names, frame_groups=frame_groups)
+
+

imgtools/modules/structureset.py

@@ -50,7 +50,6 @@ def from_dicom_rtstruct(cls, rtstruct_path: str) -> 'StructureSet':
         metadata = {}
 
         return cls(roi_points, metadata)
-        # return cls(roi_points)
 
     @property
     def roi_names(self) -> List[str]:
@@ -186,7 +185,6 @@ def to_segmentation(self, reference_image: sitk.Image,
         if isinstance(roi_names, str):
             roi_names = [roi_names]
         if isinstance(roi_names, list): # won't this always trigger after the previous?
-            print("triggered")
             labels = self._assign_labels(roi_names, roi_select_first)
         print("labels:", labels)
         all_empty = True