Write 3D Slicer image and segmentation renderings to image and video files

save_volume_rendering.py

"""
Render image, segmentation, and colorized image slices and 3D.

How to call this from command line:

    /home/jakob/dev/Slicer-5.7.0-linux-amd64/Slicer --python-script save_volume_rendering.py /path/to/output/folder mytest /path_to/ct_15mm.nii.gz /path/to/segmentation.nii.gz

"""

def captureSliceSweep(outputFilename, numberOfImages=25, lightboxColumns=5):
    """outputFilename: path to the output lightbox image file (e.g., "c:/tmp/slices.png")
    """
    import pathlib
    import ScreenCapture
    cap = ScreenCapture.ScreenCaptureLogic()

    # Capture slice sweep
    sliceScreenshotsFilenamePattern = pathlib.Path(outputFilename).parent.joinpath("slices_%04d.png")
    cap.showViewControllers(False)
    slicer.app.layoutManager().resetSliceViews()
    sliceNode = slicer.util.getNode("vtkMRMLSliceNodeRed")
    sliceOffsetMin, sliceOffsetMax = cap.getSliceOffsetRange(sliceNode)
    sliceOffsetStart = sliceOffsetMin + (sliceOffsetMax - sliceOffsetMin) * 0.05
    sliceOffsetEnd = sliceOffsetMax - (sliceOffsetMax - sliceOffsetMin) * 0.05
    cap.captureSliceSweep(
        sliceNode, sliceOffsetStart, sliceOffsetEnd, numberOfImages,
        sliceScreenshotsFilenamePattern.parent, sliceScreenshotsFilenamePattern.name,
        captureAllViews=None, transparentBackground=False)
    cap.showViewControllers(True)

    # Create lightbox image
    cap.createLightboxImage(lightboxColumns,
        sliceScreenshotsFilenamePattern.parent,
        sliceScreenshotsFilenamePattern.name,
        numberOfImages,
        outputFilename)
    cap.deleteTemporaryFiles(sliceScreenshotsFilenamePattern.parent, sliceScreenshotsFilenamePattern.name, numberOfImages)

def capture3DRotation(outputFilename, numberOfImages=25, lightboxColumns=5, videoLengthSec=5):
    """
    outputFilename: path to the output lightbox image filel gif, mp4, or png. If png then lightbox image is saved, otherwise video is created.
    """
    import pathlib
    import ScreenCapture
    cap = ScreenCapture.ScreenCaptureLogic()

    # Capture 3D rotation
    rotate3dScreenshotsFilenamePattern = pathlib.Path(outputFilename).parent.joinpath("rotate3d_%04d.png")

    if outputFilename.suffix.lower() == ".png":
        video = False
    else:
        video = True
        if outputFilename.suffix.lower() == ".gif":
            # animated GIF
            extraOptions = "-filter_complex palettegen,[v]paletteuse"
        elif outputFilename.suffix.lower() == ".mp4":
            # H264 high-quality
            extraOptions = "-codec libx264 -preset slower -crf 18 -pix_fmt yuv420p"
        else:
            raise ValueError(f"Unsupported format: {outputFilename.suffix}")

    viewNode = slicer.app.layoutManager().threeDWidget(0).mrmlViewNode()
    viewNode.SetBackgroundColor(0,0,0)
    viewNode.SetBackgroundColor2(0,0,0)
    viewNode.SetAxisLabelsVisible(False)
    viewNode.SetBoxVisible(False)
    cap.showViewControllers(False)
    slicer.app.layoutManager().resetThreeDViews()
    cap.capture3dViewRotation(viewNode, -180, 180, numberOfImages, ScreenCapture.AXIS_YAW, rotate3dScreenshotsFilenamePattern.parent, rotate3dScreenshotsFilenamePattern.name)
    cap.showViewControllers(True)

    if video:
        cap.createVideo(numberOfImages/videoLengthSec, extraOptions, rotate3dScreenshotsFilenamePattern.parent, rotate3dScreenshotsFilenamePattern.name, outputFilename)
    else:
        cap.createLightboxImage(lightboxColumns,
            rotate3dScreenshotsFilenamePattern.parent,
            rotate3dScreenshotsFilenamePattern.name,
            numberOfImages,
            outputFilename)

    cap.deleteTemporaryFiles(rotate3dScreenshotsFilenamePattern.parent, rotate3dScreenshotsFilenamePattern.name, numberOfImages)

def captureRenderings(outputPath, name, image, segmentation, width, height):
    import pathlib
    # Set up visualization
    
    modulePanel = slicer.util.findChild(slicer.util.mainWindow(), "PanelDockWidget")
    modulePanel.hide()

    import qt
    slicer.util.mainWindow().size = qt.QSize(width, height)

    slicer.util.setPythonConsoleVisible(False)
    slicer.app.layoutManager().setLayout(slicer.vtkMRMLLayoutNode.SlicerLayoutFourUpView)

    viewNode = slicer.app.layoutManager().threeDWidget(0).mrmlViewNode()
    segmentation.GetDisplayNode().SetVisibility(False)
    segmentation.CreateClosedSurfaceRepresentation()
    segmentation.GetDisplayNode().SetVisibility(True)
    segmentation.GetDisplayNode().SetOpacity3D(0.6)
    segmentation.GetDisplayNode().SetOpacity2DFill(0.1)
    vrDisplayNode = slicer.modules.volumerendering.logic().CreateDefaultVolumeRenderingNodes(image)
    vrDisplayNode.SetVisibility(True)
    slicer.app.layoutManager().resetThreeDViews()

    # Capture slices and 3D of CT
    vrDisplayNode.SetVisibility(True)
    segmentation.GetDisplayNode().SetVisibility(False)
    viewNode.SetShadowsVisibility(True)
    captureSliceSweep(pathlib.Path(outputPath).joinpath(f"{name}_ct_slices.png"))
    capture3DRotation(pathlib.Path(outputPath).joinpath(f"{name}_ct_3d.gif"))

    # Capture slices and 3D of segmentation
    vrDisplayNode.SetVisibility(False)
    segmentation.GetDisplayNode().SetVisibility(True)
    viewNode.SetShadowsVisibility(False)
    captureSliceSweep(pathlib.Path(outputPath).joinpath(f"{name}_segmentation_slices.png"))
    capture3DRotation(pathlib.Path(outputPath).joinpath(f"{name}_segmentation_3d.gif"))

    # Generate colorized volume
    from ColorizeVolume import ColorizeVolumeLogic
    colorizeLogic = ColorizeVolumeLogic()
    parameterNode = colorizeLogic.getParameterNode()
    parameterNode.inputScalarVolume = image
    parameterNode.inputSegmentation = segmentation
    parameterNode.outputRgbaVolume = colorizeLogic.AddNewOutputVolume()
    colorizeLogic.process()

    # Capture 3D of colorized volume
    vrDisplayNode.SetVisibility(False)
    segmentation.GetDisplayNode().SetVisibility(False)
    colorizeLogic.showVolumeRendering()
    viewNode.SetShadowsVisibility(True)
    capture3DRotation(pathlib.Path(outputPath).joinpath(f"{name}_colorized_ct_3d.gif"))

    modulePanel.show()


if __name__ == "__main__":
    import sys

    # Get input data
    if len(sys.argv) == 5:
        outputPath = sys.argv[1]
        name = sys.argv[2]
        image = slicer.util.loadVolume(sys.argv[3])
        segmentation = slicer.util.loadSegmentation(sys.argv[4])
    else:
        print("""Usage:              
              Slicer --python-script save_volume_rendering.py <path_to_output_folder> <name> <path_to_ct_image> <path_to_segmentation>")
              """)
        print(sys.argv)
        sys.exit(1)

    captureRenderings(outputPath, name, image, segmentation, 1000, 800)
    exit(0)