MedicalDemo4

#!/usr/bin/env python

import sys
from dataclasses import dataclass
from pathlib import Path

# noinspection PyUnresolvedReferences
import vtkmodules.vtkInteractionStyle
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingOpenGL2
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingVolumeOpenGL2
from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkCommonDataModel import vtkPiecewiseFunction
from vtkmodules.vtkCommonTransforms import vtkTransform
from vtkmodules.vtkFiltersGeneral import vtkTransformFilter
from vtkmodules.vtkFiltersModeling import vtkOutlineFilter
from vtkmodules.vtkIOImage import vtkMetaImageReader
from vtkmodules.vtkImagingCore import vtkImageChangeInformation
from vtkmodules.vtkInteractionWidgets import vtkCameraOrientationWidget
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkCamera,
    vtkColorTransferFunction,
    vtkPolyDataMapper,
    vtkProperty,
    vtkRenderWindow,
    vtkRenderWindowInteractor,
    vtkRenderer,
    vtkVolume,
    vtkVolumeProperty
)
from vtkmodules.vtkRenderingVolume import vtkFixedPointVolumeRayCastMapper


def main():
    colors = vtkNamedColors()

    file_name = get_program_parameters()

    colors.SetColor('BkgColor', 51, 77, 102, 255)

    # Create the renderer, the render window, and the interactor. The renderer
    # draws into the render window, the interactor enables mouse- and
    # keyboard-based interaction with the scene.
    ren = vtkRenderer()
    ren_win = vtkRenderWindow()
    ren_win.AddRenderer(ren)
    # Set a background color for the renderer and set the name and
    # size of the render window (expressed in pixels).
    ren = vtkRenderer(background=colors.GetColor3d('BkgColor'))
    ren_win = vtkRenderWindow(size=(640, 480),
                              window_name=f'{Path(sys.argv[0]).stem:s}')
    ren_win.AddRenderer(ren)

    iren = vtkRenderWindowInteractor()
    iren.SetRenderWindow(ren_win)
    # Since we import vtkmodules.vtkInteractionStyle we can do this
    # because vtkInteractorStyleSwitch is automatically imported:
    iren.GetInteractorStyle().SetCurrentStyleToTrackballCamera()

    # The following reader is used to read a series of 2D slices (images)
    # that compose the volume. The slice dimensions are set, and the
    # pixel spacing. The data Endianness must also be specified. The reader
    # uses the FilePrefix in combination with the slice number to construct
    # filenames using the format FilePrefix.%d. (In this case the FilePrefix
    # is the root name of the file: quarter.)
    reader = vtkMetaImageReader(file_name=file_name)
    reader.SetDataByteOrderToLittleEndian()
    reader.update()
    # Get the current physical center of the volume
    # Bounds returns: [xmin, xmax, ymin, ymax, zmin, zmax]
    bounds = reader.output.bounds
    centroid = list()
    for i in range(0, len(bounds), 2):
        centroid.append(bounds[i] + (bounds[i + 1] - bounds[i]) / 2.0)

    # Apply the shift using vtkImageChangeInformation.
    change_information = vtkImageChangeInformation()
    reader >> change_information
    # change_information.SetInputConnection(reader.GetOutputPort())
    # Center the image by shifting it by the negative center coordinates.
    change_information.SetOriginTranslation(list(map(lambda x: -x, centroid)))

    # An outline provides context around the data.
    outline_filter = vtkOutlineFilter()
    change_information >> outline_filter
    # We need to transform the outline so that it outlines the head.
    transform = vtkTransform()
    # Move the outline from the origin to the new center inside the head.
    transform.Translate(-centroid[0] * 2.0, -centroid[1] / 32.0, -centroid[2] * 2.0)
    # Apply the transform to the geometry.
    transform_filter = vtkTransformFilter(input_connection=outline_filter.output_port, transform=transform)
    transform_filter.update()

    outline_mapper = vtkPolyDataMapper()
    transform_filter >> outline_mapper
    outline_property = vtkProperty(diffuse_color=colors.GetColor3d('Black'))
    outline = vtkActor(mapper=outline_mapper, property=outline_property)
    ren.AddActor(outline)

    # The volume will be displayed by ray-cast alpha compositing.
    # A ray-cast mapper is needed to do the ray-casting.
    volume_mapper = vtkFixedPointVolumeRayCastMapper()
    change_information >> volume_mapper

    # The color transfer function maps voxel intensities to colors.
    # It is modality-specific, and often anatomy-specific as well.
    # The goal is to one color for flesh (between 500 and 1000)
    # and another color for bone (1150 and over).
    volume_color = vtkColorTransferFunction()
    volume_color.AddRGBPoint(0, 0.0, 0.0, 0.0)
    volume_color.AddRGBPoint(500, 240.0 / 255.0, 184.0 / 255.0, 160.0 / 255.0)
    volume_color.AddRGBPoint(1000, 240.0 / 255.0, 184.0 / 255.0, 160.0 / 255.0)
    volume_color.AddRGBPoint(1150, 1.0, 1.0, 240.0 / 255.0)  # Ivory

    # The opacity transfer function is used to control the opacity
    # of different tissue types.
    volume_scalar_opacity = vtkPiecewiseFunction()
    volume_scalar_opacity.AddPoint(0, 0.00)
    volume_scalar_opacity.AddPoint(500, 0.15)
    volume_scalar_opacity.AddPoint(1000, 0.15)
    volume_scalar_opacity.AddPoint(1150, 0.85)

    # The gradient opacity function is used to decrease the opacity
    # in the 'flat' regions of the volume while maintaining the opacity
    # at the boundaries between tissue types.  The gradient is measured
    # as the amount by which the intensity changes over unit distance.
    # For most medical data, the unit distance is 1mm.
    volume_gradient_opacity = vtkPiecewiseFunction()
    volume_gradient_opacity.AddPoint(0, 0.0)
    volume_gradient_opacity.AddPoint(90, 0.5)
    volume_gradient_opacity.AddPoint(100, 1.0)

    # The VolumeProperty attaches the color and opacity functions to the
    # volume, and sets other volume properties.  The interpolation should
    # be set to linear in order to do a high-quality rendering.
    # The ShadeOn option turns on directional lighting, which will usually
    # enhance the appearance of the volume and make it look more '3D'.
    # However the quality of the shading depends on how accurately the gradient
    # of the volume can be calculated, and for noisy data the gradient
    # estimation will be very poor. The impact of the shading can be
    # decreased by increasing the Ambient coefficient while decreasing
    # the Diffuse and Specular coefficient. To increase the impact
    # of shading, decrease the Ambient and increase the Diffuse and Specular.
    volume_property = vtkVolumeProperty(color=volume_color, scalar_opacity=volume_scalar_opacity,
                                        gradient_opacity=volume_gradient_opacity, shade=True,
                                        ambient=0.4, diffuse=0.6, specular=0.2)
    volume_property.SetInterpolationTypeToLinear()

    # The vtkVolume is a vtkProp3D (like a vtkActor) and controls the position
    # and orientation of the volume in world coordinates.
    volume = vtkVolume(property=volume_property, mapper=volume_mapper)

    # Finally, add the volume to the renderer
    ren.AddViewProp(volume)

    # Set up an initial view of the volume.  The focal point will be the
    # center of the volume, and the camera position will be 400mm to the
    # patient's left (which is our right).
    camera = vtkCamera(view_up=(0, 0, -1), position=(-400, 0, 0), focal_point=(0, 0, 0))
    camera.ComputeViewPlaneNormal()
    camera.Elevation(30.0)
    ren.SetActiveCamera(camera)
    camera.Dolly(1.5)
    ren.ResetCameraClippingRange()

    ren_win.Render()

    category = 'rlpais'
    cow = make_camera_orientation_widget(ren, category)
    # Enable the widget.
    cow.On()

    iren.Initialize()
    iren.Start()


def get_program_parameters():
    import argparse
    description = 'Read a volume dataset and display it via volume rendering.'
    epilogue = '''
    Derived from VTK/Examples/Cxx/Medical4.cxx
    '''
    parser = argparse.ArgumentParser(description=description, epilog=epilogue,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('filename', help='FullHead.mhd.')
    args = parser.parse_args()
    return args.filename



def set_axes_labels():
    """
    Define the axes labels.

    :return: The axes labels.
    """
    return {
        # Labels are: Anterior, Posterior, Dorsal, Ventral, Left, Right
        'apdvlr': {'+X': 'A', '-X': 'P', '+Y': 'D', '-Y': 'V', '+Z': 'L', '-Z': 'R'},
        'apdvrl': {'+X': 'A', '-X': 'P', '+Y': 'D', '-Y': 'V', '+Z': 'R', '-Z': 'L'},
        'padvlr': {'+X': 'P', '-X': 'A', '+Y': 'D', '-Y': 'V', '+Z': 'L', '-Z': 'R'},
        # Labels are: Left, Right, Superior, Inferior, Anterior, Posterior
        'lrsiap': {'+X': 'L', '-X': 'R', '+Y': 'S', '-Y': 'I', '+Z': 'A', '-Z': 'P'},
        'lrpasi': {'+X': 'L', '-X': 'R', '+Y': 'P', '-Y': 'A', '+Z': 'S', '-Z': 'I'},
        'rlpais': {'+X': 'R', '-X': 'L', '+Y': 'P', '-Y': 'A', '+Z': 'I', '-Z': 'S'},
        # Default labels
        'xyz': None
    }


def set_axes_colors():
    """
    Define the axes colors.

    :return: The axes colors.
    """
    color1 = {'+X': 'IndianRed', '-X': 'FireBrick',
              '+Y': 'LimeGreen', '-Y': 'DarkGreen',
              '+Z': 'Blue', '-Z': 'SteelBlue'}
    return {
        'apdvlr': color1,
        'apdvrl': color1,
        'padvlr': color1,
        'lrsiap': color1,
        'lrpasi': None,
        'rlpais': None,
        # Default colors.
        'xyz': None
    }


def get_axes_params():
    """
    Gather the defined axes labels and colors into a dictionary.

    :return: The dictionary of axes labels and colors.
    """

    # The keys must be the same.
    axes_labels = set_axes_labels()
    axes_colors = set_axes_colors()
    label_keys = set(axes_labels.keys())
    color_keys = set(axes_colors.keys())
    common_keys = label_keys.intersection(color_keys)
    alc = dict()
    for k in common_keys:
        alc[k] = (axes_labels[k], axes_colors[k])
    return alc


def make_camera_orientation_widget(ren, alc_key='xyz', position=3):  # , reposition):
    """
    Make a camera orientation widget for a given renderer.

    position has these values 0: LowerLeft, 1: UpperLeft, 2: LowerRight, 3: UpperRight

    :param ren: The renderer.
    :param alc_key: The key specifying the desired labels and colors for the axes.
    :param position: Position the camera orientation widget.
    :return: The camera orientation widget.
    """

    cow = vtkCameraOrientationWidget(parent_renderer=ren, enabled=True)

    rep = cow.representation

    axes_parameters = get_axes_params()

    if alc_key not in axes_parameters:
        print(
            f'Invalid key for axes labels and colors.'
            f'\nValid keys are: {sorted(axes_parameters.keys())}'
            f'\nUsing the key: xyz.')
        alc_key = 'xyz'

    alc = axes_parameters[alc_key]

    match position:
        case 0:
            rep.AnchorToLowerLeft()
        case 1:
            rep.AnchorToUpperLeft()
        case 2:
            rep.AnchorToLowerRight()
        case _:
            rep.AnchorToUpperRight()

    if not alc[0] is None:
        rep.x_plus_label_text = alc[0]['+X']
        rep.x_minus_label_text = alc[0]['-X']
        rep.y_plus_label_text = alc[0]['+Y']
        rep.y_minus_label_text = alc[0]['-Y']
        rep.z_plus_label_text = alc[0]['+Z']
        rep.z_minus_label_text = alc[0]['-Z']

    if not alc[1] is None:
        colors = vtkNamedColors()
        rep.x_axis_color = colors.GetColor3d(alc[1]['+X'])
        rep.y_axis_color = colors.GetColor3d(alc[1]['+Y'])
        rep.z_axis_color = colors.GetColor3d(alc[1]['+Z'])

    cow.SetRepresentation(rep)

    cow.Off()

    return cow



if __name__ == '__main__':
    main()