BUG: Cleaning up load of DICOM segmentation (issue QIICR#194)

- referenced master volume is set properly only if SEG will be loaded after referenced volumes has been loaded. This is not always the case though! In some cases scalar volume is loaded, then SEG and in the end the PET SUV, which causes the SEG to set the wrong master volume as master volume - using Subject Hierarchy for achieving querying for the referenced volume - adding references from ReferencedImageSequence besides ReferencedSeriesSequence
che85 · Oct 2, 2017 · 71dff6e · 71dff6e
1 parent 0af4868
commit 71dff6e
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 48 deletions.
diff --git a/DICOMPlugins/DICOMSegmentationPlugin.py b/DICOMPlugins/DICOMSegmentationPlugin.py
@@ -2,6 +2,7 @@
 import os
 import json
 import vtk
+import vtkSegmentationCorePython as vtkSegmentationCore
 import logging
 
 from base.DICOMPluginBase import DICOMPluginBase
@@ -185,7 +186,6 @@ def load(self,loadable):
                                     loadable.name,
                                     regionCode, regionCodingScheme, regionCodeMeaning,
                                     regionModCode, regionModCodingScheme, regionModCodeMeaning)
-          # end of processing a line of terminology
 
           # TODO: Create logic class that both CLI and this plugin uses so that we don't need to have temporary NRRD
           # files and labelmap nodes
@@ -194,68 +194,99 @@ def load(self,loadable):
           # load the segmentation volume file and name it for the reference series and segment color
           labelFileName = os.path.join(self.tempDir, str(segmentId) + ".nrrd")
           segmentName = seriesName + "-" + typeCodeMeaning + "-label"
-          (success, labelNode) = slicer.util.loadLabelVolume(labelFileName,
-                                                             properties={'name': segmentName},
-                                                             returnNode=True)
+          success, labelNode = slicer.util.loadLabelVolume(labelFileName,properties={'name': segmentName},
+                                                           returnNode=True)
           if not success:
             raise ValueError("{} could not be loaded into Slicer!".format(labelFileName))
-          segmentLabelNodes.append(labelNode)
-
+
           # Set terminology properties as attributes to the label node (which is a temporary node)
           #TODO: This is a quick solution, maybe there is a better one
           labelNode.SetAttribute("Terminology", segmentTerminologyTag)
           labelNode.SetAttribute("ColorR", str(rgb[0]))
           labelNode.SetAttribute("ColorG", str(rgb[1]))
           labelNode.SetAttribute("ColorB", str(rgb[2]))
 
-          # Create subject hierarchy for the loaded series
-          self.addSeriesInSubjectHierarchy(loadable, labelNode)
-
-      metaFile.close()
+          segmentLabelNodes.append(labelNode)
 
     self.cleanup()
 
-    import vtkSegmentationCorePython as vtkSegmentationCore
-    vtkSegConverter = vtkSegmentationCore.vtkSegmentationConverter
+    self._createSegmentationNode(loadable, segmentLabelNodes)
 
-    segmentationNode = slicer.vtkMRMLSegmentationNode()
-    segmentationNode.SetName(seriesName)
-    slicer.mrmlScene.AddNode(segmentationNode)
+    return True
 
-    segmentationDisplayNode = slicer.vtkMRMLSegmentationDisplayNode()
-    slicer.mrmlScene.AddNode(segmentationDisplayNode)
-    segmentationNode.SetAndObserveDisplayNodeID(segmentationDisplayNode.GetID())
+  def _createSegmentationNode(self, loadable, segmentLabelNodes):
+    segmentationNode = self._initializeSegmentation(loadable)
 
-    segmentation = vtkSegmentationCore.vtkSegmentation()
-    segmentation.SetMasterRepresentationName(vtkSegConverter.GetSegmentationBinaryLabelmapRepresentationName())
+    for segmentLabelNode in segmentLabelNodes:
+      self._importSegmentAndRemoveLabel(segmentLabelNode, segmentationNode)
 
-    segmentationNode.SetAndObserveSegmentation(segmentation)
     self.addSeriesInSubjectHierarchy(loadable, segmentationNode)
+    self._findAndSetGeometryReference(loadable.referencedSeriesUID, segmentationNode)
 
-    for segmentLabelNode in segmentLabelNodes:
-      segment = vtkSegmentationCore.vtkSegment()
-      segment.SetName(segmentLabelNode.GetName())
+  def _importSegmentAndRemoveLabel(self, segmentLabelNode, segmentationNode):
+    segmentationsLogic = slicer.modules.segmentations.logic()
+    segmentation = segmentationNode.GetSegmentation()
+    success = segmentationsLogic.ImportLabelmapToSegmentationNode(segmentLabelNode, segmentationNode)
+    if success:
+      segment = segmentation.GetNthSegment(segmentation.GetNumberOfSegments() - 1)
+      segment.SetColor([float(segmentLabelNode.GetAttribute("ColorR")),
+                        float(segmentLabelNode.GetAttribute("ColorG")),
+                        float(segmentLabelNode.GetAttribute("ColorB"))])
 
-      segmentColor = [float(segmentLabelNode.GetAttribute("ColorR")), float(segmentLabelNode.GetAttribute("ColorG")),
-                      float(segmentLabelNode.GetAttribute("ColorB"))]
-      segment.SetColor(segmentColor)
-
       segment.SetTag(vtkSegmentationCore.vtkSegment.GetTerminologyEntryTagName(),
                      segmentLabelNode.GetAttribute("Terminology"))
 
-      #TODO: when the logic class is created, this will need to be changed
-      orientedImage = slicer.vtkSlicerSegmentationsModuleLogic.CreateOrientedImageDataFromVolumeNode(segmentLabelNode)
-      segment.AddRepresentation(vtkSegConverter.GetSegmentationBinaryLabelmapRepresentationName(), orientedImage)
-      segmentation.AddSegment(segment)
+      self._removeLabelNode(segmentLabelNode)
+    return segmentation
+
+  def _removeLabelNode(self, labelNode):
+    dNode = labelNode.GetDisplayNode()
+    if dNode is not None:
+      slicer.mrmlScene.RemoveNode(dNode)
+    slicer.mrmlScene.RemoveNode(labelNode)
 
-      segmentDisplayNode = segmentLabelNode.GetDisplayNode()
-      if segmentDisplayNode is not None:
-        slicer.mrmlScene.RemoveNode(segmentDisplayNode)
-      slicer.mrmlScene.RemoveNode(segmentLabelNode)
+  def _initializeSegmentation(self, loadable):
+    segmentationNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationNode")
+    segmentationNode.SetName(self.referencedSeriesName(loadable))
 
+    segmentationDisplayNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationDisplayNode")
+    segmentationNode.SetAndObserveDisplayNodeID(segmentationDisplayNode.GetID())
+
+    vtkSegConverter = vtkSegmentationCore.vtkSegmentationConverter
+    segmentation = vtkSegmentationCore.vtkSegmentation()
+    segmentation.SetMasterRepresentationName(vtkSegConverter.GetSegmentationBinaryLabelmapRepresentationName())
     segmentation.CreateRepresentation(vtkSegConverter.GetSegmentationClosedSurfaceRepresentationName(), True)
+    segmentationNode.SetAndObserveSegmentation(segmentation)
 
-    return True
+    return segmentationNode
+
+  def _findAndSetGeometryReference(self, referencedSeriesUID, segmentationNode):
+    shn = slicer.vtkMRMLSubjectHierarchyNode.GetSubjectHierarchyNode(slicer.mrmlScene)
+    segID = shn.GetItemByDataNode(segmentationNode)
+    parentID = shn.GetItemParent(segID)
+    childIDs = vtk.vtkIdList()
+    shn.GetItemChildren(parentID, childIDs)
+    uidName = slicer.vtkMRMLSubjectHierarchyConstants.GetDICOMUIDName()
+    matches = []
+    for childID in reversed([childIDs.GetId(id) for id in range(childIDs.GetNumberOfIds()) if segID]):
+      if childID == segID:
+        continue
+      if shn.GetItemUID(childID, uidName) == referencedSeriesUID:
+        if shn.GetItemDataNode(childID):
+          matches.append(shn.GetItemDataNode(childID))
+
+    reference = None
+    if len(matches):
+      if any(x.GetAttribute("DICOM.RWV.instanceUID") is not None for x in matches):
+        for x in matches:
+          if x.GetAttribute("DICOM.RWV.instanceUID") is not None:
+            reference = x
+            break
+      else:
+        reference = matches[0]
+
+    if reference:
+      segmentationNode.SetReferenceImageGeometryParameterFromVolumeNode(reference)
 
   def examineForExport(self, subjectHierarchyItemID):
     shNode = slicer.vtkMRMLSubjectHierarchyNode.GetSubjectHierarchyNode(slicer.mrmlScene)
@@ -361,8 +392,6 @@ def createLabelNodeFromSegment(segmentationNode, segmentID):
     slicer.mrmlScene.AddNode(labelNode)
     segmentationsLogic = slicer.modules.segmentations.logic()
 
-    import vtkSegmentationCorePython as vtkSegmentationCore
-
     mergedImageData = vtkSegmentationCore.vtkOrientedImageData()
     segmentationNode.GenerateMergedLabelmapForAllSegments(mergedImageData, 0, None,
                                                           DICOMSegmentationExporter.vtkStringArrayFromList([segmentID]))
@@ -456,7 +485,7 @@ def export(self, outputDirectory, segFileName, segmentIDs=None, skipEmpty=False)
       waitCount += 1
 
     if cliNode.GetStatusString() != 'Completed':
-      raise RuntimeError("itkimage2segimage CLI did not complete cleanly")
+      raise RuntimeError("%s:\n\n%s" % (cliNode.GetStatusString(), cliNode.GetErrorText()))
 
     if not os.path.exists(segFilePath):
       raise RuntimeError("DICOM Segmentation was not created. Check Error Log for further information.")
@@ -469,7 +498,6 @@ def getNonEmptySegmentIDs(self, segmentIDs):
     return [segmentID for segmentID in segmentIDs if not self.isSegmentEmpty(segmentation.GetSegment(segmentID))]
 
   def isSegmentEmpty(self, segment):
-    import vtkSegmentationCorePython as vtkSegmentationCore
     vtkSegConverter = vtkSegmentationCore.vtkSegmentationConverter
     r = segment.GetRepresentation(vtkSegConverter.GetSegmentationBinaryLabelmapRepresentationName())
     imagescalars = r.GetPointData().GetArray("ImageScalars")

diff --git a/DICOMPlugins/base/DICOMPluginBase.py b/DICOMPlugins/base/DICOMPluginBase.py
@@ -1,4 +1,5 @@
 import slicer
+import dicom
 import logging
 from datetime import datetime
 from DICOMLib import DICOMPlugin
@@ -39,17 +40,24 @@ def cleanup(self):
   def addReferences(self, loadable):
     """Puts a list of the referenced UID into the loadable for use
     in the node if this is loaded."""
-    import dicom
     dcm = dicom.read_file(loadable.files[0])
+    loadable.referencedInstanceUIDs = []
+    self._addReferencedSeries(loadable, dcm)
+    self._addReferencedImages(loadable, dcm)
+    loadable.referencedInstanceUIDs = list(set(loadable.referencedInstanceUIDs))
 
+  def _addReferencedSeries(self, loadable, dcm):
     if hasattr(dcm, "ReferencedSeriesSequence"):
-      # look up all of the instances in the series, since segmentation frames
-      #  may be non-contiguous
       if hasattr(dcm.ReferencedSeriesSequence[0], "SeriesInstanceUID"):
-        loadable.referencedInstanceUIDs = []
         for f in slicer.dicomDatabase.filesForSeries(dcm.ReferencedSeriesSequence[0].SeriesInstanceUID):
           refDCM = dicom.read_file(f)
-          # this is a hack that should probably fixed in Slicer core - not all
-          #  of those instances are truly referenced!
           loadable.referencedInstanceUIDs.append(refDCM.SOPInstanceUID)
-          loadable.referencedSeriesUID = dcm.ReferencedSeriesSequence[0].SeriesInstanceUID
+        loadable.referencedSeriesUID = dcm.ReferencedSeriesSequence[0].SeriesInstanceUID
+
+  def _addReferencedImages(self, loadable, dcm):
+    if hasattr(dcm, "ReferencedImageSequence"):
+      for item in dcm.ReferencedImageSequence:
+        if hasattr(item, "ReferencedSOPInstanceUID"):
+          loadable.referencedInstanceUIDs.append(item.ReferencedSOPInstanceUID)
+          # TODO: what to do with the SeriesInstanceUID?
+          # refDCM = dicom.read_file(slicer.dicomDatabase.fileForInstance(item.ReferencedSOPInstanceUID))