Python script to take screenshots of the VTK error meshes from the CETUS challenge

vtk_heart.py

#!/usr/bin/python
"""
http://www.it.uu.se/edu/course/homepage/avgrafik/vt06/ComputerExercises/vtkexamples.shtml
"""

import sys
import vtk
import numpy as np
import math
from math import cos, sin
import os
from glob import glob

def vtk_basic( actors, save="", magnification=3 ):
    """
    Create a window, renderer, interactor, add the actors and start the thing
    
    Parameters
    ----------
    actors :  list of vtkActors
    
    Returns
    -------
    nothing
    """     
    
    # create a rendering window and renderer
    ren = vtk.vtkRenderer()
    renWin = vtk.vtkRenderWindow()
    renWin.AddRenderer(ren)
    renWin.SetSize(600,600)
    ren.SetBackground( 1, 1, 1)
 
    # create a renderwindowinteractor
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(renWin)

    for a in actors:
        # assign actor to the renderer
        ren.AddActor(a )

    # style = vtk.vtkInteractorStyleTerrain() 
    # iren.SetInteractorStyle( style )
    
    # render
    renWin.Render()
   
    if save:
        if not save.endswith('.png'):
            save += '.png'
        grabber = vtk.vtkWindowToImageFilter()
        grabber.SetInput( renWin )
        grabber.SetMagnification( magnification )
        grabber.Update()
        
        writer = vtk.vtkPNGWriter()
        writer.SetInput( grabber.GetOutput() )
        writer.SetFileName( save )
        writer.Write()

    else:
        # enable user interface interactor
        iren.Initialize()
        iren.Start()

def move(actor, matrix):
    transfo_mat = vtk.vtkMatrix4x4()
    for i in range(0,4):
        for j in range(0,4):
            transfo_mat.SetElement(i,j, matrix[i,j])
        
    print matrix
    print transfo_mat

    actor.SetUserMatrix(transfo_mat)

def Rx( theta ):
    theta = float(theta)/180.0*math.pi
    return np.array([[ 1,          0,           0, 0],
                     [ 0, cos(theta), -sin(theta), 0],
                     [ 0, sin(theta),  cos(theta), 0],
                     [ 0,          0,           0, 1]],dtype='float32')

def Ry( theta ):
    theta = float(theta)/180.0*math.pi
    return np.array([[  cos(theta), 0, sin(theta), 0],
                     [           0, 1,          0, 0],
                     [ -sin(theta), 0, cos(theta), 0],
                     [           0, 0,          0, 1]],dtype='float32')

def Rz( theta ):
    theta = float(theta)/180.0*math.pi
    return np.array([[ cos(theta), -sin(theta), 0, 0],
                     [ sin(theta),  cos(theta), 0, 0],
                     [          0,           0, 1, 0],
                     [          0,           0, 0, 1]],dtype='float32')

def vtk_screenshot( filename, output ):
    reader = vtk.vtkPolyDataReader()
    reader.SetFileName( filename)
    reader.Update()

    # find the range of the point scalars
    a,b = reader.GetOutput().GetPointData().GetScalars().GetRange()
    nm =  reader.GetOutput().GetPointData().GetScalars().GetName()

    # show how to print a string in python, it is similar to C-style sprintf
    print "Range of %s: %4.2f-%4.2f" %(nm,a,b)

    # transfer function (lookup table) for mapping point scalar data
    # to colors (parent class is vtkScalarsToColors)
    lut = vtk.vtkColorTransferFunction()
    lut.AddRGBPoint(a,         0.0, 0.0, 1.0)
    lut.AddRGBPoint(a+(b-a)/4, 0.0, 0.5, 0.5)
    lut.AddRGBPoint(a+(b-a)/2, 0.0, 1.0, 0.0)
    lut.AddRGBPoint(b-(b-a)/4, 0.5, 0.5, 0.0)
    lut.AddRGBPoint(b,         1.0, 0.0, 0.0)

    mapper = vtk.vtkPolyDataMapper()
    mapper.SetLookupTable(lut)
    mapper.SetScalarRange(a,b)
    mapper.SetScalarModeToUsePointData()
    mapper.SetInput( reader.GetOutput() )

    actor = vtk.vtkActor()
    actor.SetMapper(mapper)

    mat = np.eye(4)
    mat[0,0] = -1
    move( actor, np.dot(Ry(-90),Rx(-90)) )
    
    vtk_basic( [actor], save=output )

    return

all_files = glob("script/Results_1030/*")
for f in all_files:
    name = os.path.basename(f)[:-len(".vtk")]
    vtk_screenshot( f, "script/vtk_errors/"+name+".png")