Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 89 additions & 76 deletions plugins/slicer/MONAILabel/MONAILabel.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down