Skip to content

Commit

Permalink
Add automatic test and improve API of ColorizeVolume module
Browse files Browse the repository at this point in the history
  • Loading branch information
@lassoan
lassoan committed Jul 5, 2024
1 parent c8ddc75 commit 9ec7521
Showing 1 changed file with 121 additions and 54 deletions.
175 changes: 121 additions & 54 deletions ColorizeVolume/ColorizeVolume.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,44 @@ def __init__(self, parent):
This file was originally developed by Steve Pieper, Isomics and Andras Lasso, PerkLab.
"""

# Additional initialization step after application startup is complete
slicer.app.connect("startupCompleted()", registerSampleData)


#
# Register sample data sets in Sample Data module
#

def registerSampleData():
"""
Add data sets to Sample Data module.
"""
# It is always recommended to provide sample data for users to make it easy to try the module,
# but if no sample data is available then this method (and associated startupCompeted signal connection) can be removed.

import SampleData
iconsPath = os.path.join(os.path.dirname(__file__), 'Resources/Icons')

# To ensure that the source code repository remains small (can be downloaded and installed quickly)
# it is recommended to store data sets that are larger than a few MB in a Github release.

# CardiacAgatstonScoring1
SampleData.SampleDataLogic.registerCustomSampleDataSource(
# Category and sample name displayed in Sample Data module
category='Sandbox',
sampleName='CTLiverSegmentation',
# Thumbnail should have size of approximately 260x280 pixels and stored in Resources/Icons folder.
# It can be created by Screen Capture module, "Capture all views" option enabled, "Number of images" set to "Single".
thumbnailFileName=os.path.join(iconsPath, 'CTLiverSegmentation.png'),
# Download URL and target file name
uris="https://github.com/PerkLab/SlicerSandbox/releases/download/TestingData/CTLiverSegmentation.seg.nrrd",
fileNames='CTLiverSegmentation.seg.nrrd',
# Checksum to ensure file integrity. Can be computed by this command:
# import hashlib; print(hashlib.sha256(open(filename, "rb").read()).hexdigest())
checksums='SHA256:ce9a7182a666788a2556f6cf4f59ad5dadd944171cc279e80c164496729a7032',
# This node name will be used when the data set is loaded
nodeNames='CTLiverSegmentation'
)

#
# ColorizeVolumeParameterNode
Expand All @@ -61,6 +99,8 @@ class ColorizeVolumeParameterNode:
inputScalarVolume: vtkMRMLScalarVolumeNode
inputSegmentation: vtkMRMLSegmentationNode
outputRgbaVolume: vtkMRMLVectorVolumeNode
# 3-element list of floats
backgroundColorRgb: list[float] = [0.1, 0.1, 0.1]
softEdgeThicknessVoxel: Annotated[float, WithinRange(0, 8)] = 4.0
colorBleedThicknessVoxel: Annotated[float, WithinRange(0, 8)] = 1.0
backgroundOpacityPercent: Annotated[float, WithinRange(0, 100)] = 20
Expand Down Expand Up @@ -115,6 +155,7 @@ def setup(self) -> None:
# Create logic class. Logic implements all computations that should be possible to run
# in batch mode, without a graphical user interface.
self.logic = ColorizeVolumeLogic()
self.logic.logCallback = self.log

# Connections

Expand All @@ -130,6 +171,7 @@ def setup(self) -> None:
self.ui.resetVolumeRenderingSettingsButton.connect('clicked(bool)', self.onResetVolumeRenderingSettingsButton)
self.ui.volumeRenderingSettingsButton.connect('clicked(bool)', self.onVolumeRenderingSettingsButton)
self.ui.resetToDefaultsButton.connect('clicked(bool)', self.onResetToDefaultsButton)
self.ui.backgroundColorPickerButton.connect('colorChanged(QColor)', self.onBackgroundColorSelected)

self.ui.volumeRenderingLevelWidget.connect('valueChanged(double)', self.onUpdateVolumeRenderingTransferFunction)
self.ui.volumeRenderingWindowWidget.connect('valueChanged(double)', self.onUpdateVolumeRenderingTransferFunction)
Expand Down Expand Up @@ -222,6 +264,10 @@ def _checkCanApply(self, caller=None, event=None) -> None:
else:
self.ui.applyButton.toolTip = "Select input volume and segmentation"
self.ui.applyButton.enabled = False
# Parameter node modified, make sure color selector is updated
# (need to update manually because color picker is not supported by parameter node)
import qt
self.ui.backgroundColorPickerButton.color = qt.QColor.fromRgbF(*self._parameterNode.backgroundColorRgb)

def onApplyButton(self) -> None:
"""
Expand All @@ -236,25 +282,19 @@ def onApplyButton(self) -> None:
with slicer.util.tryWithErrorDisplay("Failed to compute results.", waitCursor=True):

if not self._parameterNode.outputRgbaVolume:
self._parameterNode.outputRgbaVolume = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLVectorVolumeNode", f"{self._parameterNode.inputScalarVolume.GetName()} colored")

# Parameter node does not support color selector
backgroundColorRgba = [
self.ui.backgroundColorPickerButton.color.redF(),
self.ui.backgroundColorPickerButton.color.greenF(),
self.ui.backgroundColorPickerButton.color.blueF(),
self._parameterNode.backgroundOpacityPercent / 100.0]
self._parameterNode.outputRgbaVolume = self.logic.AddNewOutputVolume(f"{self._parameterNode.inputScalarVolume.GetName()} colored")

# Compute output

import time
startTime = time.time()

self.logic.process(
self.logic._process(
self._parameterNode.inputScalarVolume,
self._parameterNode.inputSegmentation,
self._parameterNode.outputRgbaVolume,
backgroundColorRgba,
self._parameterNode.backgroundColorRgb,
self._parameterNode.backgroundOpacityPercent,
self._parameterNode.colorBleedThicknessVoxel,
self._parameterNode.softEdgeThicknessVoxel,
sequenceBrowserNode)
Expand Down Expand Up @@ -293,6 +333,9 @@ def onOutputRgbaVolumeSelected(self) -> None:
return
self.ui.volumePropertyNodeWidget.setMRMLVolumePropertyNode(volumeRenderingPropertyNode)

def onBackgroundColorSelected(self, color) -> None:
self._parameterNode.backgroundColorRgb = [color.redF(), color.greenF(), color.blueF()]

def onResetToDefaultsButton(self):
for paramName in ['softEdgeThicknessVoxel', 'colorBleedThicknessVoxel', 'backgroundOpacityPercent']:
self.logic.getParameterNode().setValue(paramName, self.logic.getParameterNode().default(paramName).value)
Expand All @@ -302,6 +345,11 @@ def onUpdateVolumeRenderingTransferFunction(self):
# We call the logic update via a timer to give time for the parameter node to get updated.
import qt
qt.QTimer.singleShot(0, self.logic.updateVolumeRenderingOpacityTransferFunctions)

def log(self, message):
slicer.util.showStatusMessage(message, 1000)
slicer.app.processEvents()


#
# ColorizeVolumeLogic
Expand All @@ -322,16 +370,29 @@ def __init__(self) -> None:
Called when the logic class is instantiated. Can be used for initializing member variables.
"""
ScriptedLoadableModuleLogic.__init__(self)
self.logCallback = None

def log(self, message):
if self.logCallback:
self.logCallback(message)
else:
logging.info(message)

def getParameterNode(self):
return ColorizeVolumeParameterNode(super().getParameterNode())

def processVolume(
def AddNewOutputVolume(self, nodeName=None):
outputRgbaVolume = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLVectorVolumeNode", nodeName if nodeName else "")
outputRgbaVolume.SetVoxelVectorType(slicer.vtkMRMLVolumeNode.VoxelVectorTypeColorRGBA)
return outputRgbaVolume

def _processVolume(
self,
volumeNode,
segmentationNode,
outputRgbaVolume,
backgroundColorRgba,
backgroundColorRgb,
backgroundColorOpacityPercent,
colorBleedThicknessVoxel,
softEdgeThicknessVoxel,
):
Expand All @@ -341,7 +402,8 @@ def processVolume(
:param volumeNode: volume to be thresholded
:param segmentationNode: segmentation to be used for coloring
:param outputRgbaVolume: colorized RGBA volume
:param backgroundColorRgba: color and opacity of voxels that are not segmented (RGBA)
:param backgroundColorRgb: color and opacity of voxels that are not segmented (RGBA)
:param backgroundColorOpacityPercent: opacity of voxels that are not segmented (0-100)
:oaram colorBleedThicknessVoxel: how far color bleeds out (in voxels)
:param softEdgeThicknessVoxel: edge smoothing thickness (in voxels)
"""
Expand All @@ -350,25 +412,23 @@ def processVolume(
import vtk
import vtk.util.numpy_support
segmentIds = segmentationNode.GetDisplayNode().GetVisibleSegmentIDs()
slicer.util.showStatusMessage("Exporting segments...", 1000)
slicer.app.processEvents()
self.log("Exporting segments...")
labelmapVolumeNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLLabelMapVolumeNode", "__temp__")
if not slicer.modules.segmentations.logic().ExportSegmentsToLabelmapNode(segmentationNode, segmentIds, labelmapVolumeNode, volumeNode):
raise RuntimeError("Export of segment failed.")
slicer.app.processEvents()
colorTableNode = labelmapVolumeNode.GetDisplayNode().GetColorNode()

# Background color
colorTableNode.SetColor(0, *backgroundColorRgba)
colorTableNode.SetColor(0, *backgroundColorRgb, backgroundColorOpacityPercent / 100.0)
for segmentIndex, segmentId in enumerate(segmentIds):
segment = segmentationNode.GetSegmentation().GetSegment(segmentId)
color = segment.GetColor()
opacity = segmentationNode.GetDisplayNode().GetSegmentOpacity3D(segmentId)
colorTableNode.SetColor(segmentIndex + 1, *color, opacity)

# Dilate labelmap to avoid edge artifacts
slicer.util.showStatusMessage(f"Dilating segments...")
slicer.app.processEvents()
self.log("Dilating segments...")
dilate = vtkAddon.vtkImageLabelDilate3D()
dilate.SetInputData(labelmapVolumeNode.GetImageData())
dilationKernelSize = int(colorBleedThicknessVoxel + 0.5) * 2 + 1
Expand All @@ -377,8 +437,7 @@ def processVolume(
dilate.Update()
labelImage = dilate.GetOutput()

slicer.util.showStatusMessage(f"Generating colorized volume...")
slicer.app.processEvents()
self.log("Generating colorized volume...")

mapToRGB = vtk.vtkImageMapToColors()
mapToRGB.ReleaseDataFlagOn()
Expand Down Expand Up @@ -451,11 +510,12 @@ def processVolume(
slicer.mrmlScene.RemoveNode(colorTableNode)
slicer.mrmlScene.RemoveNode(labelmapVolumeNode)

def process(self,
def _process(self,
inputScalarVolume: vtkMRMLScalarVolumeNode,
inputSegmentation: vtkMRMLSegmentationNode,
outputRgbaVolume: vtkMRMLVectorVolumeNode,
backgroundColorRgba: list,
backgroundColorRgb: list,
backgroundOpacityPercent: float,
colorBleedThicknessVoxel: float=1.5,
softEdgeThicknessVoxel: float=1.5,
sequenceBrowserNode: Optional[vtkMRMLSequenceBrowserNode]=None,
Expand All @@ -466,14 +526,19 @@ def process(self,
:param inputScalarVolume: volume to be thresholded
:param inputSegmentation: segmentation to be used for coloring
:param outputRgbaVolume: colorized RGBA volume
:param backgroundColorRgba: color and opacity of voxels that are not segmented (RGBA)
:param backgroundColorRgb: color of voxels that are not segmented (RGB)
:param backgroundOpacityPercent: opacity of voxels that are not segmented (0-100)
:oaram colorBleedThicknessVoxel: how far color bleeds out (in voxels)
:param softEdgeThicknessVoxel: edge smoothing thickness (in voxels)
:param sequenceBrowserNode: if set to a browser node then all items of the input volume sequence will be processed
"""

if not inputScalarVolume or not inputSegmentation or not outputRgbaVolume:
raise ValueError("Input or output volume is invalid")
if not inputScalarVolume:
raise ValueError("Input scalar volume is invalid")
if not inputSegmentation:
raise ValueError("Input segmentation is invalid")
if not outputRgbaVolume:
raise ValueError("Output RGBA volume is invalid")

import numpy as np
import time
Expand All @@ -489,11 +554,12 @@ def process(self,

if inputScalarVolumeSequence is None:
# Colorize a single volume
self.processVolume(
self._processVolume(
inputScalarVolume,
inputSegmentation,
outputRgbaVolume,
backgroundColorRgba,
backgroundColorRgb,
backgroundOpacityPercent,
colorBleedThicknessVoxel,
softEdgeThicknessVoxel,
)
Expand All @@ -516,11 +582,12 @@ def process(self,
numberOfItems = sequenceBrowserNode.GetNumberOfItems()
for i in range(numberOfItems):
logging.info(f"Colorizing item {i+1}/{numberOfItems} of sequence")
self.processVolume(
ColorizeVolumeLogic._processVolume(
inputScalarVolume,
inputSegmentation,
outputRgbaVolume,
backgroundColorRgba,
backgroundColorRgb,
backgroundOpacityPercent,
colorBleedThicknessVoxel,
softEdgeThicknessVoxel,
)
Expand All @@ -530,8 +597,18 @@ def process(self,
stopTime = time.time()
logging.info(f'Processing completed in {stopTime-startTime:.2f} seconds')

slicer.util.showStatusMessage("Processing completed.", 1000)
self.log("Processing completed.")

def process(self):
parameterNode = self.getParameterNode()
self._process(
parameterNode.inputScalarVolume,
parameterNode.inputSegmentation,
parameterNode.outputRgbaVolume,
parameterNode.backgroundColorRgb,
parameterNode.backgroundOpacityPercent,
parameterNode.colorBleedThicknessVoxel,
parameterNode.softEdgeThicknessVoxel)

def showVolumeRendering(self, resetSettings = True) -> None:
"""
Expand Down Expand Up @@ -654,34 +731,24 @@ def test_ColorizeVolume1(self):

self.delayDisplay("Starting the test")

# Get/create input data

# Get input data
import SampleData
registerSampleData()
inputScalarVolume = SampleData.downloadSample('ColorizeVolume1')
inputScalarVolume = SampleData.downloadSample('CTLiver')
inputSegmentation = SampleData.downloadSample('CTLiverSegmentation')
self.delayDisplay('Loaded test data set')

inputScalarRange = inputScalarVolume.GetImageData().GetScalarRange()
self.assertEqual(inputScalarRange[0], 0)
self.assertEqual(inputScalarRange[1], 695)

outputRgbaVolume = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLScalarVolumeNode")
threshold = 100

# Test the module logic

# Generate colorized volume
logic = ColorizeVolumeLogic()

# Test algorithm with non-inverted threshold
logic.process(inputScalarVolume, outputRgbaVolume, threshold, True)
outputScalarRange = outputRgbaVolume.GetImageData().GetScalarRange()
self.assertEqual(outputScalarRange[0], inputScalarRange[0])
self.assertEqual(outputScalarRange[1], threshold)

# Test algorithm with inverted threshold
logic.process(inputScalarVolume, outputRgbaVolume, threshold, False)
outputScalarRange = outputRgbaVolume.GetImageData().GetScalarRange()
self.assertEqual(outputScalarRange[0], inputScalarRange[0])
self.assertEqual(outputScalarRange[1], inputScalarRange[1])
logic.logCallback = lambda msg: self.delayDisplay(msg)
parameterNode = logic.getParameterNode()
parameterNode.inputScalarVolume = inputScalarVolume
parameterNode.inputSegmentation = inputSegmentation
parameterNode.outputRgbaVolume = logic.AddNewOutputVolume()
logic.process()

# Show volume rendering
logic.showVolumeRendering()
slicer.app.layoutManager().resetThreeDViews()

self.delayDisplay('Test passed')

0 comments on commit 9ec7521

Please sign in to comment.