TissueLens

#!/usr/bin/env python
import sys
from pathlib import Path

# noinspection PyUnresolvedReferences
import vtkmodules.vtkInteractionStyle
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingOpenGL2
from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkCommonCore import (
    vtkLookupTable
)
from vtkmodules.vtkCommonDataModel import vtkSphere
from vtkmodules.vtkFiltersCore import (
    vtkFlyingEdges3D,
    vtkMarchingCubes,
    vtkProbeFilter
)
from vtkmodules.vtkFiltersGeneral import vtkClipDataSet
from vtkmodules.vtkFiltersSources import vtkSphereSource
from vtkmodules.vtkIOImage import vtkMetaImageReader
from vtkmodules.vtkInteractionWidgets import vtkCameraOrientationWidget
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkCamera,
    vtkDataSetMapper,
    vtkProperty,
    vtkRenderWindow,
    vtkRenderWindowInteractor,
    vtkRenderer
)


def main():
    colors = vtkNamedColors()

    file_name, use_flying_edges = get_program_parameters()

    colors.SetColor('SkinColor', 240, 184, 160, 255)
    colors.SetColor('BackfaceColor', 255, 229, 200, 255)
    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 data within the render window.
    #
    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()

    # Read the volume data
    reader = vtkMetaImageReader(file_name=file_name)
    reader.update()

    # An isosurface, or contour value of 500 is known to correspond to the
    # skin of the patient.
    if use_flying_edges:
        try:
            skin_extractor = vtkFlyingEdges3D()
        except AttributeError:
            skin_extractor = vtkMarchingCubes()
    else:
        skin_extractor = vtkMarchingCubes()
    skin_extractor.input_connection = reader.output_port
    skin_extractor.SetValue(0, 500)

    clip_center = (-70, 60, -10)

    # Define a spherical clip function to clip the isosurface
    clip_function = vtkSphere(radius=50, center=clip_center)

    # Clip the isosurface with a sphere
    skin_clip = vtkClipDataSet(input_connection=skin_extractor.output_port, clip_function=clip_function, value=0,
                               generate_clip_scalars=True)
    skin_clip.update()

    skin_mapper = vtkDataSetMapper(input_connection=skin_clip.output_port, scalar_visibility=False)

    skin_prop = vtkProperty(diffuse_color=colors.GetColor3d('SkinColor'))
    back_prop = vtkProperty(diffuse_color=colors.GetColor3d('BackfaceColor'))

    skin = vtkActor(mapper=skin_mapper, property=skin_prop, backface_property=back_prop)

    # Define a model for the "lens". Its geometry matches the implicit
    # sphere used to clip the isosurface
    lens_model = vtkSphereSource(radius=50, center=clip_center, phi_resolution=201, theta_resolution=101)

    # Sample the input volume with the lens model geometry
    lens_probe = vtkProbeFilter(input_connection=lens_model.output_port, source_connection=reader.output_port)

    # Clip the lens data with the isosurface value
    lens_clip = vtkClipDataSet(input_connection=lens_probe.output_port, value=500, generate_clip_scalars=False)
    lens_clip.update()

    # Define a suitable grayscale lut
    bw_lut = vtkLookupTable()
    bw_lut.SetTableRange(0, 2048)
    bw_lut.SetSaturationRange(0, 0)
    bw_lut.SetHueRange(0, 0)
    bw_lut.SetValueRange(0.2, 1)
    bw_lut.Build()

    lens_mapper = vtkDataSetMapper(input_connection=lens_clip.output_port, scalar_range=lens_clip.output.scalar_range,
                                   lookup_table=bw_lut)

    lens = vtkActor(mapper=lens_mapper)
    # lens.SetMapper(lens_mapper)

    # It is convenient to create an initial view of the data. The FocalPoint
    # and Position form a vector direction. Later on (ResetCamera() method)
    # this vector is used to position the camera to look at the data in
    # this direction.
    camera = vtkCamera(view_up=(0, 0, 1), position=(0, -1, 0), focal_point=(0, 0, 0))
    camera.ComputeViewPlaneNormal()
    camera.Azimuth(30.0)
    camera.Elevation(30.0)

    # Actors are added to the renderer. An initial camera view is created.
    # The Dolly() method moves the camera towards the FocalPoint,
    # thereby enlarging the image.
    ren.AddActor(lens)
    ren.AddActor(skin)
    ren.SetActiveCamera(camera)
    ren.ResetCamera()
    camera.Dolly(1.5)

    # Note that when camera movement occurs (as it does in the Dolly()
    # method), the clipping planes often need adjusting. Clipping planes
    # consist of two planes: near and far along the view direction. The
    # near plane clips out objects in front of the plane the far plane
    # clips out objects behind the plane. This way only what is drawn
    # between the planes is actually rendered.
    ren.ResetCameraClippingRange()

    ren_win.Render()

    cow = make_camera_orientation_widget(ren)
    cow.On()

    category = 'lrpasi'
    cow1 = make_camera_orientation_widget(ren, category, 1)
    cow1.On()

    # Initialize the event loop and then start it.
    iren.Initialize()
    iren.Start()


def get_program_parameters():
    import argparse
    description = ('Generate a "magnifying lens" or "window" effect, allowing'
                   ' users to peek inside a 3D medical dataset')
    epilogue = '''
    '''
    parser = argparse.ArgumentParser(description=description, epilog=epilogue,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('filename', help='FullHead.mhd.')
    parser.add_argument('-m', '--marching_cubes', action='store_false',
                        help='Use Marching Cubes instead of Flying Edges.')
    args = parser.parse_args()
    return args.filename, args.marching_cubes


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()