Logo Search packages:      
Sourcecode: mayavi2 version File versions  Download package

volume.py

00001 """The Volume module visualizes scalar fields using volumetric
visualization techniques.  This supports ImageData and
UnstructuredGrid data.  It also supports the FixedPointRenderer for
ImageData.  However, the performance is slow so your best bet is
probably with the ImageData based renderers.

"""
# Author: Prabhu Ramachandran <prabhu@aero.iitb.ac.in>
# Copyright (c) 2006, Enthought, Inc.
# License: BSD Style.

# Standard imports
from math import cos, sqrt, pi
from vtk.util import vtkConstants

# Enthought library imports.
from enthought.traits.api import Instance, Property, List, ReadOnly, \
     Str, Button, Tuple
from enthought.traits.ui.api import View, Group, Item, InstanceEditor, CustomEditor
from enthought.tvtk.api import tvtk
from enthought.tvtk.util.gradient_editor import hsva_to_rgba, GradientTable
from enthought.tvtk.util.ctf import save_ctfs, load_ctfs, \
     rescale_ctfs, set_lut, PiecewiseFunction, ColorTransferFunction
from enthought.tvtk.util.wx_gradient_editor import gradient_editor_factory
from enthought.persistence import state_pickler

# Local imports
from enthought.mayavi.core.module import Module
from enthought.mayavi.core.common import error
from enthought.mayavi.core.traits import DEnum
from enthought.mayavi.core.lut_manager import LUTManager

######################################################################
# Utility functions.
######################################################################
00036 def is_volume_pro_available():
    """Returns `True` if there is a volume pro card available.
    """
    try:
        map = tvtk.VolumeProMapper()
    except AttributeError:
        return False
    else:
        return map.number_of_boards > 0

def find_volume_mappers():
    res = []
    for name in dir(tvtk):
        if 'Volume' in name and 'Mapper' in name and 'OpenGL' not in name:
            try:
                klass = getattr(tvtk, name)
                inst = klass()
            except TypeError:
                pass
            else:
                res.append(name)
    ignores = ['VolumeTextureMapper3D', 'VolumeProMapper']
    for name in ignores:
        if name in res:
            res.remove(name)
    return res

00063 def default_OTF(x1, x2):
    """Creates a default opacity transfer function.
    """
    maxs = max(x1, x2)
    mins = min(x1, x2)
    otf = PiecewiseFunction()
    otf.add_point(mins, 0.0)
    otf.add_point(maxs, 0.2)
    return otf

00073 def make_CTF(x1, x2, hue_range=(2.0/3.0, 0.0),
             sat_range=(1.0, 1.0), val_range=(1.0, 1.0),
             n=10, mode='sqrt'):
    """Creates a CTF as per given arguments.  Lets one set a hue,
    saturation and value range.  The mode can be one of 'sqrt', or
    'linear'.  The idea of this function is to create a CTF that is
    similar to the LUT being used.  's_curve' is not implemented.
    Patches welcome.
    """
    maxs = max(x1, x2)
    mins = min(x1, x2)
    ds = maxs - mins
    dhue = hue_range[1] - hue_range[0]
    dsat = sat_range[1] - sat_range[0]
    dval = val_range[1] - val_range[0]
    ctf = ColorTransferFunction()
    if mode == 'sqrt':
        for i in range(n+1):
            # Generate x in [0, 1]
            x = 0.5*(1.0  + cos((n-i)*pi/n)) # Chebyshev nodes.
            h = hue_range[0] + dhue*x
            s = sat_range[0] + dsat*x
            v = val_range[0] + dval*x
            r, g, b, a = [sqrt(c) for c in hsva_to_rgba(h, s, v, 1.0)]
            ctf.add_rgb_point(mins+x*ds, r, g, b)
    elif mode == 'linear':
        for i in range(n+1):
            # Generate x in [0, 1]
            x = float(i)/n # Uniform nodes.
            h = hue_range[0] + dhue*x
            s = sat_range[0] + dsat*x
            v = val_range[0] + dval*x
            r, g, b, a = hsva_to_rgba(h, s, v, 1.0)
            ctf.add_rgb_point(mins+x*ds, r, g, b)      
    return ctf
    

00110 def default_CTF(x1, x2):
    """Creates a default RGB color transfer function.  In this case we
    default to a red-blue one with the 'sqrt' mode.
    """
    return make_CTF(x1, x2,
                    hue_range=(2.0/3.0, 0.0),
                    sat_range=(1.0, 1.0),
                    val_range=(1.0, 1.0),
                    n=10,
                    mode='sqrt')


00122 def load_volume_prop_from_grad(grad_file_name, volume_prop,
                               scalar_range=(0, 255)):
    """Load a ``*.grad`` file (*grad_file_name*) and set the given volume
    property (*volume_prop*) given the *scalar_range*.
    """
    gt = GradientTable(300)
    gt.load(grad_file_name)
    gt.store_to_vtk_volume_prop(volume_prop, scalar_range)


00132 def save_volume_prop_to_grad(volume_prop, grad_file_name):
    """Save the given volume property (*volume_prop*) to a ``*.grad`` file
    given as *grad_file_name*.
    """
    gt = GradientTable(300)
    gt.load_from_vtk_volume_prop(volume_prop)
    gt.save(grad_file_name)


######################################################################
# `VolumeLutManager` class.
######################################################################
00144 class VolumeLUTManager(LUTManager):
    """Just has a different view than the LUTManager.
    """
    view = View(Group(Item(name='show_scalar_bar'),
                      Item(name='number_of_labels'),
                      Item(name='shadow'),
                      Item(name='use_default_name'),
                      Item(name='data_name'),
                      label='Scalar Bar',
                      ),
                Group(Item(name='_title_text_property',
                           style='custom',
                           resizable=True),
                      show_labels=False,
                      label='Title'),
                Group(Item(name='_label_text_property',
                           style='custom',
                           resizable=True),
                      show_labels=False,
                      label='Labels'),
                resizable=True
                )
    

######################################################################
# `Volume` class.
######################################################################
00171 class Volume(Module):
    """The Volume module visualizes scalar fields using volumetric
    visualization techniques.  This supports ImageData and
    UnstructuredGrid data.  It also supports the FixedPointRenderer
    for ImageData.  However, the performance is slow so your best bet
    is probably with the ImageData based renderers.
    """

    # The version of this class.  Used for persistence.
    __version__ = 0

    volume_mapper_type = DEnum(values_name='_mapper_types',
                               desc='volume mapper to use')

    ray_cast_function_type = DEnum(values_name='_ray_cast_functions',
                                   desc='Ray cast function to use')
    
    volume = ReadOnly

    volume_mapper = Property

    volume_property = Property

    ray_cast_function = Property

    lut_manager = Instance(VolumeLUTManager, args=(), allow_none=False)

    ########################################
    # View related code.

    update_ctf = Button('Update CTF')
    
    view = View(Group(Item(name='_volume_property', style='custom',
                           editor=CustomEditor(gradient_editor_factory),
                           resizable=True),
                      Item(name='update_ctf'),
                      label='CTF',
                      show_labels=False),
                Group(Item(name='volume_mapper_type'),
                      Group(Item(name='_volume_mapper',
                                 style='custom',
                                 resizable=True),
                            show_labels=False
                            ),
                      Item(name='ray_cast_function_type'),
                      Group(Item(name='_ray_cast_function',
                                 enabled_when='len(_ray_cast_functions) > 0',
                                 style='custom',
                                 resizable=True),
                            show_labels=False),
                      label='Mapper',
                      ),
                Group(Item(name='_volume_property', style='custom',
                           resizable=True),
                      label='Property',
                      show_labels=False),                
                Group(Item(name='volume', style='custom',
                           editor=InstanceEditor(),
                           resizable=True),
                      label='Volume',
                      show_labels=False),
                Group(Item(name='lut_manager', style='custom',
                           resizable=True),
                      label='Legend',
                      show_labels=False),                
                resizable=True
                )

    ########################################
    # Private traits
    _volume_mapper = Instance(tvtk.AbstractVolumeMapper)
    _volume_property = Instance(tvtk.VolumeProperty)
    _ray_cast_function = Instance(tvtk.Object)

    _mapper_types = List(Str, ['TextureMapper2D', 'RayCastMapper', ])

    _available_mapper_types = List(Str)

    _ray_cast_functions = List(Str)

    current_range = Tuple

    # The color transfer function.
    _ctf = Instance(ColorTransferFunction)
    # The opacity values.
    _otf = Instance(PiecewiseFunction)

    ######################################################################
    # `object` interface
    ######################################################################
    def __get_pure_state__(self):
        d = super(Volume, self).__get_pure_state__()
        d['ctf_state'] = save_ctfs(self._volume_property)
        for name in ('current_range', '_ctf', '_otf'):
            d.pop(name, None)
        return d

    def __set_pure_state__(self, state):
        self.volume_mapper_type = state['_volume_mapper_type']
        state_pickler.set_state(self, state, ignore=['ctf_state'])
        ctf_state = state['ctf_state']
        ctf, otf = load_ctfs(ctf_state, self._volume_property)
        self._ctf = ctf
        self._otf = otf
        self._update_ctf_fired()
        
    ######################################################################
    # `Module` interface
    ######################################################################
    def start(self):
        super(Volume, self).start()
        self.lut_manager.start()

    def stop(self):
        super(Volume, self).stop()
        self.lut_manager.stop()
    
00288     def setup_pipeline(self):
        """Override this method so that it *creates* the tvtk
        pipeline.
        """
        v = self.volume = tvtk.Volume()
        vp = self._volume_property = tvtk.VolumeProperty()

        self._ctf = ctf = default_CTF(0, 255)
        self._otf = otf = default_OTF(0, 255)
        vp.set_color(ctf)
        vp.set_scalar_opacity(otf)
        vp.shade = True
        vp.interpolation_type = 'linear'
        v.property = vp

        v.on_trait_change(self.render)
        vp.on_trait_change(self.render)

        available_mappers = find_volume_mappers()
        if is_volume_pro_available():
            self._mapper_types.append('VolumeProMapper')
            available_mappers.append('VolumeProMapper')

        self._available_mapper_types = available_mappers
        if 'FixedPointVolumeRayCastMapper' in available_mappers:
            self._mapper_types.append('FixedPointVolumeRayCastMapper')
        
        self.actors.append(v)
    
00317     def update_pipeline(self):
        """Override this method so that it *updates* the tvtk pipeline
        when data upstream is known to have changed.

        This method is invoked (automatically) when any of the inputs
        sends a `pipeline_changed` event.
        """
        mm = self.module_manager
        if mm is None:
            return

        input = mm.source.outputs[0]

        ug = hasattr(tvtk, 'UnstructuredGridVolumeMapper')
        if ug:
            if not input.is_a('vtkImageData') \
                   and not input.is_a('vtkUnstructuredGrid'):
                error('Volume rendering only works with '\
                      'StructuredPoints/ImageData/UnstructuredGrid datasets')
                return
        elif not input.is_a('vtkImageData'):
            error('Volume rendering only works with '\
                  'StructuredPoints/ImageData datasets')
            return
    
        self._setup_mapper_types()
        self._setup_current_range()
        self._volume_mapper_type_changed(self.volume_mapper_type)
        self._update_ctf_fired()
        self.pipeline_changed = True

00348     def update_data(self):
        """Override this method so that it flushes the vtk pipeline if
        that is necessary.

        This method is invoked (automatically) when any of the inputs
        sends a `data_changed` event.
        """
        self._setup_mapper_types()
        self._setup_current_range()
        self._update_ctf_fired()
        self.data_changed = True
    
    ######################################################################
    # Non-public methods.
    ######################################################################
00363     def _setup_mapper_types(self):
        """Sets up the mapper based on input data types.
        """
        input = self.module_manager.source.outputs[0]
        if input.is_a('vtkUnstructuredGrid'):
            if hasattr(tvtk, 'UnstructuredGridVolumeMapper'):
                check = ['UnstructuredGridVolumeZSweepMapper',
                         'UnstructuredGridVolumeRayCastMapper',
                         ]
                mapper_types = []
                for mapper in check:
                    if mapper in self._available_mapper_types:
                        mapper_types.append(mapper)
                if len(mapper_types) == 0:
                    mapper_types = ['']
                self._mapper_types = mapper_types
                return
        else:
            if input.point_data.scalars.data_type not in \
               [vtkConstants.VTK_UNSIGNED_CHAR,
                vtkConstants.VTK_UNSIGNED_SHORT]:
                if 'FixedPointVolumeRayCastMapper' \
                       in self._available_mapper_types:
                    self._mapper_types = ['FixedPointVolumeRayCastMapper']
                else:
                    error('Available volume mappers only work with \
                    unsigned_char or unsigned_short datatypes')
            else:
                mapper_types = ['TextureMapper2D', 'RayCastMapper']
                check = ['FixedPointVolumeRayCastMapper',
                         'VolumeProMapper'
                         ]
                for mapper in check:
                    if mapper in self._available_mapper_types:
                        mapper_types.append(mapper)
                self._mapper_types = mapper_types
        
    def _setup_current_range(self):
        mm = self.module_manager
        # Set the default name and range for our lut.
        lm = self.lut_manager
        slm = mm.scalar_lut_manager 
        lm.set(default_data_name=slm.default_data_name,
               default_data_range=slm.default_data_range)
        
        # Set the current range.
        input = mm.source.outputs[0]
        sc = input.point_data.scalars
        if sc is not None:
            rng = sc.range
        else:
            error('No scalars in input data!')
            rng = (0, 255)

        if self.current_range != rng:
            self.current_range = rng
            
    def _get_volume_mapper(self):
        return self._volume_mapper

    def _get_volume_property(self):
        return self._volume_property
    
    def _get_ray_cast_function(self):
        return self._ray_cast_function

    def _volume_mapper_type_changed(self, value):
        mm = self.module_manager
        if mm is None:
            return

        old_vm = self._volume_mapper
        if old_vm is not None:
            old_vm.on_trait_change(self.render, remove=True)

        if value == 'RayCastMapper':
            new_vm = tvtk.VolumeRayCastMapper()
            self._volume_mapper = new_vm
            self._ray_cast_functions = ['RayCastCompositeFunction',
                                        'RayCastMIPFunction',
                                        'RayCastIsosurfaceFunction']
            new_vm.volume_ray_cast_function = tvtk.VolumeRayCastCompositeFunction()
        elif value == 'TextureMapper2D':
            new_vm = tvtk.VolumeTextureMapper2D()
            self._volume_mapper = new_vm
            self._ray_cast_functions = ['']
        elif value == 'VolumeProMapper':
            new_vm = tvtk.VolumeProMapper()
            self._volume_mapper = new_vm
            self._ray_cast_functions = ['']
        elif value == 'FixedPointVolumeRayCastMapper':            
            new_vm = tvtk.FixedPointVolumeRayCastMapper()
            self._volume_mapper = new_vm
            self._ray_cast_functions = ['']
        elif value == 'UnstructuredGridVolumeRayCastMapper':
            new_vm = tvtk.UnstructuredGridVolumeRayCastMapper()
            self._volume_mapper = new_vm
            self._ray_cast_functions = ['']
        elif value == 'UnstructuredGridVolumeZSweepMapper':
            new_vm = tvtk.UnstructuredGridVolumeZSweepMapper()
            self._volume_mapper = new_vm
            self._ray_cast_functions = ['']

        new_vm.input = mm.source.outputs[0]
        self.volume.mapper = new_vm
        new_vm.on_trait_change(self.render)
        
    def _update_ctf_fired(self):
        set_lut(self.lut_manager.lut, self._volume_property)
        self.render()

    def _current_range_changed(self, old, new):
        rescale_ctfs(self._volume_property, new)
        self.render()
        
    def _ray_cast_function_type_changed(self, old, new):
        rcf = self.ray_cast_function
        if len(old) > 0:
            rcf.on_trait_change(self.render, remove=True)

        if len(new) > 0:
            new_rcf = getattr(tvtk, 'Volume%s'%new)()
            new_rcf.on_trait_change(self.render)
            self._volume_mapper.volume_ray_cast_function = new_rcf
            self._ray_cast_function = new_rcf
        else:
            self._ray_cast_function = None

        self.render()

    def _scene_changed(self, old, new):        
        super(Volume, self)._scene_changed(old, new)
        self.lut_manager.scene = new

Generated by  Doxygen 1.6.0   Back to index