From 07051c9dd8dfd36708665e728136f2f919d2f8b8 Mon Sep 17 00:00:00 2001 From: Janis Vahldiek Date: Thu, 5 May 2022 10:17:19 +0200 Subject: [PATCH] add support for .seg.nrrd Signed-off-by: Janis Vahldiek --- plugins/slicer/MONAILabel/MONAILabel.py | 165 +++++++++++++----------- 1 file changed, 89 insertions(+), 76 deletions(-) diff --git a/plugins/slicer/MONAILabel/MONAILabel.py b/plugins/slicer/MONAILabel/MONAILabel.py index ea152f49d..ae972eee5 100644 --- a/plugins/slicer/MONAILabel/MONAILabel.py +++ b/plugins/slicer/MONAILabel/MONAILabel.py @@ -1578,85 +1578,98 @@ def updateSegmentationMask(self, in_file, labels, sliceIndex=None, freeze=None): segmentation.AddEmptySegment(label, label, self.getLabelColor(label)) return True - labels = [label for label in labels if label != "background"] - print(f"Update Segmentation Mask using Labels: {labels}") - - # segmentId, segment = self.currentSegment() - labelImage = sitk.ReadImage(in_file) - labelmapVolumeNode = sitkUtils.PushVolumeToSlicer(labelImage, None, className="vtkMRMLLabelMapVolumeNode") - - existing_label_ids = {} - for label in labels: - id = segmentation.GetSegmentIdBySegmentName(label) - if id: - existing_label_ids[label] = id - - freeze = [freeze] if freeze and isinstance(freeze, str) else freeze - print(f"Import only Freezed label: {freeze}") - - numberOfExistingSegments = segmentation.GetNumberOfSegments() - slicer.modules.segmentations.logic().ImportLabelmapToSegmentationNode(labelmapVolumeNode, segmentationNode) - slicer.mrmlScene.RemoveNode(labelmapVolumeNode) - - numberOfAddedSegments = segmentation.GetNumberOfSegments() - numberOfExistingSegments - logging.debug(f"Adding {numberOfAddedSegments} segments") - - addedSegmentIds = [ - segmentation.GetNthSegmentID(numberOfExistingSegments + i) for i in range(numberOfAddedSegments) - ] - for i, segmentId in enumerate(addedSegmentIds): - segment = segmentation.GetSegment(segmentId) - print(f"Setting new segmentation with id: {segmentId} => {segment.GetName()}") - - label = labels[i] if i < len(labels) else f"unknown {i}" - # segment.SetName(label) - # segment.SetColor(self.getLabelColor(label)) - - if freeze and label not in freeze: - print(f"Discard label update for: {label}") - elif label in existing_label_ids: - segmentEditorWidget = slicer.modules.segmenteditor.widgetRepresentation().self().editor - segmentEditorWidget.setSegmentationNode(segmentationNode) - segmentEditorWidget.setMasterVolumeNode(self._volumeNode) - segmentEditorWidget.setCurrentSegmentID(existing_label_ids[label]) - - effect = segmentEditorWidget.effectByName("Logical operators") - labelmap = slicer.vtkOrientedImageData() - segmentationNode.GetBinaryLabelmapRepresentation(segmentId, labelmap) - - if sliceIndex: - selectedSegmentLabelmap = effect.selectedSegmentLabelmap() - dims = selectedSegmentLabelmap.GetDimensions() - count = 0 - for x in range(dims[0]): - for y in range(dims[1]): - if selectedSegmentLabelmap.GetScalarComponentAsDouble(x, y, sliceIndex, 0): - count = count + 1 - selectedSegmentLabelmap.SetScalarComponentFromDouble(x, y, sliceIndex, 0, 0) - - logging.debug(f"Total Non Zero: {count}") - - # Clear the Slice - if count: + if in_file.endswith(".seg.nrrd") and self.file_ext == ".seg.nrrd": + source_node = slicer.modules.segmentations.logic().LoadSegmentationFromFile(in_file, False) + destination_node = segmentationNode + destination_segmentations = destination_node.GetSegmentation() + source_segmentations = source_node.GetSegmentation() + + destination_segmentations.DeepCopy(source_segmentations) + + if self._volumeNode: + destination_node.SetReferenceImageGeometryParameterFromVolumeNode(self._volumeNode) + + slicer.mrmlScene.RemoveNode(source_node) + else: + labels = [label for label in labels if label != "background"] + print(f"Update Segmentation Mask using Labels: {labels}") + + # segmentId, segment = self.currentSegment() + labelImage = sitk.ReadImage(in_file) + labelmapVolumeNode = sitkUtils.PushVolumeToSlicer(labelImage, None, className="vtkMRMLLabelMapVolumeNode") + + existing_label_ids = {} + for label in labels: + id = segmentation.GetSegmentIdBySegmentName(label) + if id: + existing_label_ids[label] = id + + freeze = [freeze] if freeze and isinstance(freeze, str) else freeze + print(f"Import only Freezed label: {freeze}") + + numberOfExistingSegments = segmentation.GetNumberOfSegments() + slicer.modules.segmentations.logic().ImportLabelmapToSegmentationNode(labelmapVolumeNode, segmentationNode) + slicer.mrmlScene.RemoveNode(labelmapVolumeNode) + + numberOfAddedSegments = segmentation.GetNumberOfSegments() - numberOfExistingSegments + logging.debug(f"Adding {numberOfAddedSegments} segments") + + addedSegmentIds = [ + segmentation.GetNthSegmentID(numberOfExistingSegments + i) for i in range(numberOfAddedSegments) + ] + for i, segmentId in enumerate(addedSegmentIds): + segment = segmentation.GetSegment(segmentId) + print(f"Setting new segmentation with id: {segmentId} => {segment.GetName()}") + + label = labels[i] if i < len(labels) else f"unknown {i}" + # segment.SetName(label) + # segment.SetColor(self.getLabelColor(label)) + + if freeze and label not in freeze: + print(f"Discard label update for: {label}") + elif label in existing_label_ids: + segmentEditorWidget = slicer.modules.segmenteditor.widgetRepresentation().self().editor + segmentEditorWidget.setSegmentationNode(segmentationNode) + segmentEditorWidget.setMasterVolumeNode(self._volumeNode) + segmentEditorWidget.setCurrentSegmentID(existing_label_ids[label]) + + effect = segmentEditorWidget.effectByName("Logical operators") + labelmap = slicer.vtkOrientedImageData() + segmentationNode.GetBinaryLabelmapRepresentation(segmentId, labelmap) + + if sliceIndex: + selectedSegmentLabelmap = effect.selectedSegmentLabelmap() + dims = selectedSegmentLabelmap.GetDimensions() + count = 0 + for x in range(dims[0]): + for y in range(dims[1]): + if selectedSegmentLabelmap.GetScalarComponentAsDouble(x, y, sliceIndex, 0): + count = count + 1 + selectedSegmentLabelmap.SetScalarComponentFromDouble(x, y, sliceIndex, 0, 0) + + logging.debug(f"Total Non Zero: {count}") + + # Clear the Slice + if count: + effect.modifySelectedSegmentByLabelmap( + selectedSegmentLabelmap, slicer.qSlicerSegmentEditorAbstractEffect.ModificationModeSet + ) + + # Union label map effect.modifySelectedSegmentByLabelmap( - selectedSegmentLabelmap, slicer.qSlicerSegmentEditorAbstractEffect.ModificationModeSet + labelmap, slicer.qSlicerSegmentEditorAbstractEffect.ModificationModeAdd + ) + else: + # adding bypass masking to not overwrite other layers, + # needed for preserving scribbles during updates + # help from: https://github.com/Slicer/Slicer/blob/master + # /Modules/Loadable/Segmentations/EditorEffects/Python/SegmentEditorLogicalEffect.py + bypassMask = True + effect.modifySelectedSegmentByLabelmap( + labelmap, slicer.qSlicerSegmentEditorAbstractEffect.ModificationModeSet, bypassMask ) - # Union label map - effect.modifySelectedSegmentByLabelmap( - labelmap, slicer.qSlicerSegmentEditorAbstractEffect.ModificationModeAdd - ) - else: - # adding bypass masking to not overwrite other layers, - # needed for preserving scribbles during updates - # help from: https://github.com/Slicer/Slicer/blob/master - # /Modules/Loadable/Segmentations/EditorEffects/Python/SegmentEditorLogicalEffect.py - bypassMask = True - effect.modifySelectedSegmentByLabelmap( - labelmap, slicer.qSlicerSegmentEditorAbstractEffect.ModificationModeSet, bypassMask - ) - - segmentationNode.RemoveSegment(segmentId) + segmentationNode.RemoveSegment(segmentId) self.showSegmentationsIn3D() logging.info(f"Time consumed by updateSegmentationMask: {time.time() - start:3.1f}")