Skip to content

Commit

Permalink
Merge pull request #270 from fedorov/merged-segments
Browse files Browse the repository at this point in the history 
allow import of multiple segments from a single ITK file
  • Loading branch information
@fedorov
fedorov authored Nov 16, 2023
2 parents 8a6e612 + 562d680 commit 7f00b64
Showing 1 changed file with 62 additions and 43 deletions.
105 changes: 62 additions & 43 deletions DICOMPlugins/DICOMSegmentationPlugin.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@ def load(self,loadable):
parameters = {
"inputSEGFileName": segFileName,
"outputDirName": self.tempDir,
"mergeSegments": True,
}

try:
Expand All @@ -114,13 +115,14 @@ def load(self,loadable):
return False

cliNode = None

cliNode = slicer.cli.run(segimage2itkimage, cliNode, parameters, wait_for_completion=True)
if cliNode.GetStatusString() != 'Completed':
logging.error('SEG2NRRD did not complete successfully, unable to load DICOM Segmentation')
self.cleanup()
return False

numberOfSegments = len(glob.glob(os.path.join(self.tempDir,'*.nrrd')))
numberOfSegmentations = len(glob.glob(os.path.join(self.tempDir,'*.nrrd')))

# resize the color table to include the segments plus 0 for the background

Expand All @@ -141,18 +143,32 @@ def load(self,loadable):

with open(metaFileName) as metaFile:
data = json.load(metaFile)
logging.debug('number of segmentation files = ' + str(numberOfSegments))
if numberOfSegments != len(data["segmentAttributes"]):
logging.debug('Loaded segmentation metadata from ' + metaFileName)

logging.debug('number of segmentation files = ' + str(numberOfSegmentations))
if numberOfSegmentations != len(data["segmentAttributes"]):
logging.error('Loading failed: Inconsistent number of segments in the descriptor file and on disk')
return
for segmentAttributes in data["segmentAttributes"]:
# TODO: only handles the first item of lists

for segmentationId,segmentAttributes in enumerate(data["segmentAttributes"]):

# 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
# if not hasattr(slicer.modules, 'segmentations'):

labelFileName = os.path.join(self.tempDir, str(segmentationId+1) + ".nrrd")
# The temporary folder contains 1.nrrd, 2.nrrd,... files. We must specify singleFile=True to ensure
# they are not loaded as an image stack (could happen if each image file has a single slice only).
labelNode = slicer.util.loadLabelVolume(labelFileName, {'singleFile': True})

labelNode.labelAttributes = []

for segment in segmentAttributes:
try:
rgb255 = segment["recommendedDisplayRGBValue"]
rgb = [float(c) / 255. for c in rgb255]
except KeyError:
rgb = (0., 0., 0.)
rgb = (150., 150., 0.)

segmentId = segment["labelID"]

Expand Down Expand Up @@ -180,14 +196,6 @@ def load(self,loadable):
regionCode, regionCodingScheme, regionCodeMeaning,
regionModCode, regionModCodingScheme, regionModCodeMeaning)

# 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
# if not hasattr(slicer.modules, 'segmentations'):

labelFileName = os.path.join(self.tempDir, str(segmentId) + ".nrrd")
# The temporary folder contains 1.nrrd, 2.nrrd,... files. We must specify singleFile=True to ensure
# they are not loaded as an image stack (could happen if each image file has a single slice only).
labelNode = slicer.util.loadLabelVolume(labelFileName, {'singleFile': True})

# Set terminology properties as attributes to the label node (which is a temporary node)
segmentNameAutoGenerated = False # automatically generated from terminology
Expand All @@ -198,19 +206,21 @@ def load(self,loadable):
else:
segmentName = typeCodeMeaning
segmentNameAutoGenerated = True
labelNode.SetAttribute("Name", segmentName)
labelNode.SetAttribute("NameAutoGenerated", "1" if segmentNameAutoGenerated else "0")
labelNode.SetAttribute("Description", segment["SegmentDescription"])
labelNode.SetAttribute("Terminology", segmentTerminologyTag)
labelNode.SetAttribute("ColorR", str(rgb[0]))
labelNode.SetAttribute("ColorG", str(rgb[1]))
labelNode.SetAttribute("ColorB", str(rgb[2]))
if "SegmentAlgorithmType" in segment:
labelNode.SetAttribute("DICOM.SegmentAlgorithmType", segment["SegmentAlgorithmType"])
if "SegmentAlgorithmName" in segment:
labelNode.SetAttribute("DICOM.SegmentAlgorithmName", segment["SegmentAlgorithmName"])

segmentLabelNodes.append(labelNode)

labelAttributes = {}
labelAttributes["Name"] = segmentName
labelAttributes["NameAutoGenerated"] = segmentNameAutoGenerated
labelAttributes["Description"] = segment["SegmentDescription"]
labelAttributes["Terminology"] = segmentTerminologyTag
labelAttributes["ColorR"] = rgb[0]
labelAttributes["ColorG"] = rgb[1]
labelAttributes["ColorB"] = rgb[2]
labelAttributes["DICOM.SegmentAlgorithmType"] = segment["SegmentAlgorithmType"] if "SegmentAlgorithmType" in segment else None
labelAttributes["DICOM.SegmentAlgorithmName"] = segment["SegmentAlgorithmName"] if "SegmentAlgorithmName" in segment else None

labelNode.labelAttributes.append(labelAttributes)

segmentLabelNodes.append(labelNode)

self.cleanup()

Expand All @@ -233,28 +243,37 @@ def _importSegmentAndRemoveLabel(self, segmentLabelNode, segmentationNode):
segmentation = segmentationNode.GetSegmentation()
numberOfSegmentsBeforeImport = segmentation.GetNumberOfSegments()
success = segmentationsLogic.ImportLabelmapToSegmentationNode(segmentLabelNode, segmentationNode)

if not success:
logging.error("Failed to import segment from labelmap!")
return

if segmentation.GetNumberOfSegments() == 0:
logging.warning("Empty segment loaded from DICOM SEG!")
if segmentation.GetNumberOfSegments() - numberOfSegmentsBeforeImport > 1:
logging.warning("Multiple segments were loaded from DICOM SEG labelmap. Only one label was expected.")
if success and segmentation.GetNumberOfSegments()>0:
segment = segmentation.GetNthSegment(segmentation.GetNumberOfSegments() - 1)
segment.SetName(segmentLabelNode.GetAttribute("Name"))
segment.SetNameAutoGenerated(bool(int(segmentLabelNode.GetAttribute("NameAutoGenerated"))))
segment.SetTag("Description", segmentLabelNode.GetAttribute("Description"))
segment.SetColor([float(segmentLabelNode.GetAttribute("ColorR")),
float(segmentLabelNode.GetAttribute("ColorG")),
float(segmentLabelNode.GetAttribute("ColorB"))])

thisLabelSegmentID = 0
for segmentId in range(numberOfSegmentsBeforeImport, segmentation.GetNumberOfSegments()):
logging.info(f"Setting attributes for segment {segmentId} ...")

segment = segmentation.GetNthSegment(segmentId)
segment.SetName(segmentLabelNode.labelAttributes[thisLabelSegmentID]["Name"])
segment.SetNameAutoGenerated(segmentLabelNode.labelAttributes[thisLabelSegmentID]["NameAutoGenerated"])
segment.SetTag("Description", segmentLabelNode.labelAttributes[thisLabelSegmentID]["Description"])
segment.SetColor([float(segmentLabelNode.labelAttributes[thisLabelSegmentID]["ColorR"]),
float(segmentLabelNode.labelAttributes[thisLabelSegmentID]["ColorG"]),
float(segmentLabelNode.labelAttributes[thisLabelSegmentID]["ColorB"])])
segment.SetTag(vtkSegmentationCore.vtkSegment.GetTerminologyEntryTagName(),
segmentLabelNode.GetAttribute("Terminology"))
algorithmName = segmentLabelNode.GetAttribute("DICOM.SegmentAlgorithmName")
if algorithmName:
segmentLabelNode.labelAttributes[thisLabelSegmentID]["Terminology"])
algorithmName = segmentLabelNode.labelAttributes[thisLabelSegmentID]["DICOM.SegmentAlgorithmName"]
if algorithmName is not None:
segment.SetTag("DICOM.SegmentAlgorithmName", algorithmName)
algorithmType = segmentLabelNode.GetAttribute("DICOM.SegmentAlgorithmType")
if algorithmType:
algorithmType = segmentLabelNode.labelAttributes[thisLabelSegmentID]["DICOM.SegmentAlgorithmType"]
if algorithmType is not None:
segment.SetTag("DICOM.SegmentAlgorithmType", algorithmType)
thisLabelSegmentID += 1

self._removeLabelNode(segmentLabelNode)

self._removeLabelNode(segmentLabelNode)
return segmentation

def _removeLabelNode(self, labelNode):
Expand Down

0 comments on commit 7f00b64

Please sign in to comment.