Merge pull request #50 from prabhuramachandran/qt_gradient_editor

NEW: Adding a Qt version of the gradient editor.

mayavi/modules/volume.py

@@ -16,9 +16,10 @@
 # Enthought library imports.
 from traits.api import Instance, Property, List, ReadOnly, \
      Str, Button, Tuple
-from traitsui.api import View, Group, Item, InstanceEditor, CustomEditor
+from traitsui.api import View, Group, Item, InstanceEditor
 from tvtk.api import tvtk
 from tvtk.util.gradient_editor import hsva_to_rgba, GradientTable
+from tvtk.util.traitsui_gradient_editor import VolumePropertyEditor
 from tvtk.util.ctf import save_ctfs, load_ctfs, \
      rescale_ctfs, set_lut, PiecewiseFunction, ColorTransferFunction
 from apptools.persistence import state_pickler
@@ -33,14 +34,6 @@
 ######################################################################
 # Utility functions.
 ######################################################################
-def gradient_editor_factory(wx_parent, trait_editor):
-    """A simple wrapper to the wx specific function to avoid any UI
-    toolkit imports.
-    """
-    from tvtk.util import wx_gradient_editor as wxge
-    return wxge.gradient_editor_factory(wx_parent, trait_editor)
-
-
 def is_volume_pro_available():
     """Returns `True` if there is a volume pro card available.
     """
@@ -220,7 +213,7 @@ class Volume(Module):
     update_ctf = Button('Update CTF')
 
     view = View(Group(Item(name='_volume_property', style='custom',
-                           editor=CustomEditor(gradient_editor_factory),
+                           editor=VolumePropertyEditor,
                            resizable=True),
                       Item(name='update_ctf'),
                       label='CTF',

tvtk/util/gradient_editor.py

@@ -8,7 +8,7 @@
 This code was originally written by Gerald Knizia <cgk.d@gmx.net> and
 later modified by Prabhu Ramachandran for tvtk and MayaVi2.
 
-Copyright (c) 2005-2006, Gerald Knizia and Prabhu Ramachandran
+Copyright (c) 2005-2013, Gerald Knizia and Prabhu Ramachandran
 """
 
 from os.path import splitext
@@ -931,3 +931,338 @@ def load(self, file_name):
         self.sort_control_points()
         #self.scaling_parameters_changed()
         self.update()
+
+
+
+##########################################################################
+# `ChannelBase` class.
+##########################################################################
+class ChannelBase(object):
+    def __init__(self, function_control, name, rgb_color,
+                 channel_index, channel_mode):
+        """arguments documented in function body"""
+        self.control = function_control  #owning function control
+        self.name = name #'r','g','b','h','s','v' or 'a'
+        self.rgb_color = rgb_color
+        # ^-- string containing a tk color value with which to
+        # paint this channel
+        self.index = channel_index #0: r or h, 1: g or s, 2: b or v, 3: a
+        self.mode = channel_mode #'hsv' or 'rgb'
+
+    def get_value(self, color):
+        """Return height value of the current channel for the given color.
+        Range: 0..1"""
+        if ( self.mode == 'hsv' ):
+            return color.get_hsva()[self.index]
+        else:
+            return color.get_rgba()[self.index]
+
+    def get_value_index(self, color):
+        """Return height index of channel value of Color.
+        Range: [1..ControlHeight]"""
+        return int( 1+(self.control.height-1)*(1.0 - self.get_value(color)) )
+
+    def get_index_value(self, y):
+        """Get value in [0..1] of height index y"""
+        return min(1.0, max(0.0, 1.0 - float(y)/(self.control.height-1)))
+
+    def set_value( self, color, new_value_on_this_channel ):
+        """Color will be modified: NewValue.. will be set to the color
+        channel that ``*self`` represents."""
+        if ( self.mode == 'hsv' ):
+            hsva = [color.get_hsva()[0], color.get_hsva()[1],
+                    color.get_hsva()[2], color.get_hsva()[3] ]
+            hsva[self.index] = new_value_on_this_channel
+            if ( hsva[0] >= 1.0 - 1e-5 ):
+                # hack to make sure hue does not jump back to 0.0
+                # when it should be at 1.0 (rgb <-> hsv xform not
+                # invertible there)
+                hsva[0] = 1.0 - 1e-5
+            color.set_hsva(hsva[0],hsva[1],hsva[2],hsva[3])
+        else:
+            rgba = [color.get_rgba()[0], color.get_rgba()[1],
+                    color.get_rgba()[2], color.get_rgba()[3] ]
+            rgba[self.index] = new_value_on_this_channel
+            color.set_rgba(rgba[0],rgba[1],rgba[2],rgba[3])
+
+    def set_value_index( self, color, y ):
+        """Color will be modified: the value assigned to the height index
+        y will be set to the color channel of Color ``*self`` represents."""
+        self.set_value( color, self.get_index_value(y) )
+
+    def get_pos_index(self,f):
+        """Return x-index for gradient position f in [0..1]"""
+        return int(f*(self.control.width-1))
+
+    def get_index_pos(self,idx):
+        """Return gradient position f in [0..1] for x-index Idx in
+        [0..ControlWidth-1]"""
+        return (1.0*idx)/(self.control.width-1)
+
+    def paint(self, painter):
+        """Paint current channel into Canvas (a canvas of a function control
+        object).
+        
+        This should be overridden to do the actual painting.
+
+        """
+        raise NotImplementedError
+
+##########################################################################
+# `FunctionControl` class.
+##########################################################################
+class FunctionControl(object):
+    """Widget which displays a rectangular regions on which hue, sat, val
+    or rgb values can be modified. An function control can have one or more
+    attached color channels."""
+
+    # Radius around a control point center in which we'd still count a
+    # click as "clicked the control point"
+    control_pt_click_tolerance = 4
+
+    ChannelFactory = ChannelBase
+
+    def __init__(self, master, gradient_table, color_space, width, height):
+        """Initialize a function control widget on tkframe master.
+
+        Parameters:
+        -----------
+
+        master: The master widget.  Note that this widget *must* have
+             the methods specified in the `AbstractGradientEditorWidget`
+             interface.
+
+        on_table_changed: Callback function taking a bool argument of meaning
+            'FinalUpdate'. FinalUpdate is true if a control point is dropped,
+            created or removed and false if the update is due to a control point
+            currently beeing dragged (but not yet dropped)
+
+        color_space: String which specifies the channels painted on this control.
+             May be any combination of h,s,v,r,g,b,a in which each channel
+             occurs only once.
+
+        set_status_text: a callback used to set the status text
+             when using the editor.
+        """
+        self.text_map = {'r': 'RED', 'g': 'GREEN', 'b': 'BLUE',
+                         'h': 'HUE', 's': 'SATURATION', 'v': 'VALUE',
+                         'a': 'ALPHA'}
+        self.master = master
+        self.table = gradient_table
+        self.gradient_table = gradient_table
+        self.width = width
+        self.height = height
+
+        self.channels = []
+
+        # add the channels
+        Channel = self.ChannelFactory
+        for c in color_space:
+            if c == 'r':
+                self.channels += [Channel(self, "r", (255,0,0), 0, 'rgb' )]
+            elif c == 'g':
+                self.channels += [Channel(self, "g", (0,255,0), 1, 'rgb' )]
+            elif c == 'b':
+                self.channels += [Channel(self, "b", (0,0,255), 2, 'rgb' )]
+            elif c == 'h':
+                self.channels += [Channel(self, "h", (255,0,0), 0, 'hsv' )]
+            elif c == 's':
+                self.channels += [Channel(self, "s", (0,255,0), 1, 'hsv' )]
+            elif c == 'v':
+                self.channels += [Channel(self, "v", (0,0,255), 2, 'hsv' )]
+            elif c == 'a':
+                self.channels += [Channel(self, "a", (0,0,0), 3, 'hsv' )]
+
+        # generate a list of channels on which markers should
+        # be bound if moved on the current channel. since we interpolate
+        # the colors in hsv space, changing the r, g or b coordinates
+        # explicitely means that h, s and v all have to be fixed.
+        self.active_channels_string = ""
+        for channel in self.channels:
+            self.active_channels_string += channel.name
+        if ( ( 'r' in color_space ) or ( 'g' in color_space ) or ( 'b' in color_space ) ):
+            for c in "hsv":
+                if ( not ( c in self.active_channels_string ) ):
+                    self.active_channels_string += c
+        if ( color_space == 'a' ):
+            # alpha channels actually independent of all other channels.
+            self.active_channels_string = 'a'
+
+        # need to set to "None" initially or event handlers get confused.
+        self.cur_drag = None #<- [channel,control_point] while something is dragged.
+
+    def find_control_point(self, x, y):
+        """Check if a control point lies near (x,y) or near x if y == None.
+        returns [channel, control point] if found, None otherwise"""
+        for channel in self.channels:
+            for control_point in self.table.control_points:
+                # take into account only control points which are
+                # actually active for the current channel
+                if ( not ( channel.name in control_point.active_channels ) ):
+                    continue
+                point_x = channel.get_pos_index( control_point.pos )
+                point_y = channel.get_value_index( control_point.color )
+                y_ = y
+                if ( None == y_ ):
+                    y_ = point_y
+                if ( (point_x-x)**2 + (point_y-y_)**2 <= self.control_pt_click_tolerance**2 ):
+                    return [channel, control_point]
+        return None
+
+    def table_config_changed(self, final_update):
+        """Called internally in the control if the configuration of the attached
+        gradient table has changed due to actions of this control.
+
+        Forwards the update/change notice."""
+        self.table.update()
+        self.master.on_gradient_table_changed(final_update)
+
+    ######################################################################
+    # Toolkit specific event methods.
+    # Look at wx_gradient_editor.py and qt_gradient_editor.py to see
+    # the methods that are necessary.  
+    ######################################################################
+
+
+##########################################################################
+# `AbstractGradientEditor` interface.
+##########################################################################
+class AbstractGradientEditor(object):
+    def on_gradient_table_changed(self, final_update):
+        """ Update the gradient table and vtk lookuptable."""
+        raise NotImplementedError
+
+    def set_status_text(self, msg):
+        """Set the status on the status widget if you have one."""
+        raise NotImplementedError
+
+    def get_table_range(self):
+        """Return the CTF or LUT's scalar range."""
+        raise NotImplementedError
+
+
+##########################################################################
+# `GradientEditorWidget` interface.
+##########################################################################
+class GradientEditorWidget(AbstractGradientEditor):
+    """A Gradient Editor widget that can be used anywhere.
+    """
+    def __init__(self, master, vtk_table, on_change_color_table=None,
+                 colors=None):
+        """
+
+        Parameters:
+        -----------
+
+        vtk_table : the `tvtk.LookupTable` or `tvtk.VolumeProperty` object
+                    to set.
+
+        on_change_color_table : A callback called when the color table
+                                changes.
+
+        colors : list of 'rgb', 'hsv', 'h', 's', 'v', 'a'
+                 (Default : ['rgb', 'hsv', 'a'])
+
+                 'rgb' creates one panel to edit Red, Green and Blue
+                 colors.
+
+                 'hsv' creates one panel to edit Hue, Saturation and
+                 Value.
+
+                 'h', 's', 'v', 'r', 'g', 'b', 'a' separately
+                 specified creates different panels for each.
+        """
+        if colors is None:
+            colors = ['rgb', 'hsv', 'a']
+        self.colors = colors
+        self.gradient_preview_width = 300
+        self.gradient_preview_height = 50
+        self.channel_function_width = self.gradient_preview_width
+        self.channel_function_height = 80
+        self.gradient_table = GradientTable(self.gradient_preview_width)
+        self.vtk_color_table = vtk_table
+        if isinstance(vtk_table, tvtk.LookupTable):
+            self.vtk_table_is_lut = True
+        else:
+            # This is a tvtk.VolumeProperty
+            self.vtk_table_is_lut = False
+            # Initialize the editor with the volume property.
+            self.gradient_table.load_from_vtk_volume_prop(vtk_table)
+
+        self.on_change_color_table = on_change_color_table
+
+        # Add the function controls:
+        self.function_controls = []
+
+        self.tooltip_text = 'Left click: move control points\n'\
+                            'Right click: add/remove control points'
+        editor_data = {'rgb': ('', 'RGB'),
+                            'hsv': ('Hue: Red; Saturation: Green; '\
+                               'Value: Blue\n',
+                               'HSV'
+                               ),
+                       'h': ('', 'HUE'),
+                       's': ('', 'SAT'),
+                       'v': ('', 'VAL'),
+                       'r': ('', 'RED'),
+                       'g': ('', 'GREEN'),
+                       'b': ('', 'BLUE'),
+                       'a': ('', 'ALPHA'),
+                       }
+        self.editor_data = editor_data
+
+    ######################################################################
+    # `GradientEditorWidget` interface.
+    ######################################################################
+    def set_status_text(self, msg):
+        raise NotImplementedError
+
+    def on_gradient_table_changed(self, final_update ):
+        """ Update the gradient table and vtk lookuptable..."""
+        # update all function controls.
+        for control in self.function_controls:
+            control.update()
+        # repaint the gradient display or the external windows only
+        # when the instant*** options are set or when the update was final.
+        #if final_update or ( 1 == self.show_instant_gradients.get() ):
+        if True:
+            self.gradient_control.update()
+
+        #if final_update or ( 1 == self.show_instant_feedback.get() ):
+        if final_update:
+            vtk_table = self.vtk_color_table
+            if self.vtk_table_is_lut:
+                self.gradient_table.store_to_vtk_lookup_table(vtk_table)
+            else:
+                rng = self.get_table_range()
+                self.gradient_table.store_to_vtk_volume_prop(vtk_table, rng)
+
+            cb = self.on_change_color_table
+            if cb is not None:
+                cb()
+
+    def get_table_range(self):
+        vtk_table = self.vtk_color_table
+        if self.vtk_table_is_lut:
+            return vtk_table.table_range
+        else:
+            return vtk_table.get_scalar_opacity().range
+
+    def load(self, file_name):
+        """Set the state of the color table using the given file.
+        """
+        if len(file_name) == 0:
+            return
+        self.gradient_table.load(file_name)
+        self.on_gradient_table_changed(final_update = True)
+
+    def save(self, file_name):
+        """Store the color table to the given file.  This actually
+        generates 3 files, a '.grad', a '.lut' file and a '.jpg' file.
+        The .lut file can be used to setup a lookup table.  The .grad
+        file is used to set the state of the gradient table and the
+        JPG file is an image of the how the lut will look.
+        """
+        if len(file_name) == 0:
+            return
+        self.gradient_table.save(file_name)