<s> """<STR_LIT>""" <EOL> from traits . api import List <EOL> from tvtk . api import tvtk <EOL> def _err_msg ( obj , cls_name ) : <EOL> return '<STR_LIT>' '<STR_LIT>' % ( cls_name , str ( obj ) ) <EOL> def save_ctfs ( volume_property ) : <EOL> """<STR_LIT>""" <EOL> vp = volume_property <EOL> ctf = vp . rgb_transfer_function <EOL> otf = vp . get_scalar_opacity ( ) <EOL> s1 , s2 = ctf . range <EOL> nc = ctf . size <EOL> rgb = [ ] <EOL> if hasattr ( ctf , '<STR_LIT>' ) : <EOL> for i in range ( nc ) : <EOL> x = ctf . nodes [ i ] <EOL> r , g , b = ctf . get_color ( x ) <EOL> rgb . append ( [ x , r , g , b ] ) <EOL> elif hasattr ( ctf , '<STR_LIT>' ) : <EOL> val = [ <NUM_LIT:0> ] * <NUM_LIT:6> <EOL> for i in range ( nc ) : <EOL> ctf . get_node_value ( i , val ) <EOL> rgb . append ( val [ : <NUM_LIT:4> ] ) <EOL> else : <EOL> raise TypeError ( _err_msg ( ctf , '<STR_LIT>' ) ) <EOL> na = otf . size <EOL> a = [ ] <EOL> if hasattr ( otf , '<STR_LIT>' ) : <EOL> for i in range ( na ) : <EOL> x = otf . nodes [ i ] <EOL> val = otf . get_value ( x ) <EOL> a . append ( [ x , val ] ) <EOL> elif hasattr ( otf , '<STR_LIT>' ) : <EOL> val = [ <NUM_LIT:0> ] * <NUM_LIT:4> <EOL> for i in range ( na ) : <EOL> otf . get_node_value ( i , val ) <EOL> a . append ( val [ : <NUM_LIT:2> ] ) <EOL> else : <EOL> raise TypeError ( _err_msg ( otf , '<STR_LIT>' ) ) <EOL> return { '<STR_LIT>' : ( s1 , s2 ) , '<STR_LIT>' : rgb , '<STR_LIT>' : a } <EOL> def load_ctfs ( saved_data , volume_property ) : <EOL> """<STR_LIT>""" <EOL> rgb = saved_data [ '<STR_LIT>' ] <EOL> a = saved_data [ '<STR_LIT>' ] <EOL> new_ctf = True <EOL> ctf = volume_property . rgb_transfer_function <EOL> if isinstance ( ctf , ColorTransferFunction ) : <EOL> new_ctf = False <EOL> ctf . remove_all_points ( ) <EOL> else : <EOL> ctf = ColorTransferFunction ( ) <EOL> nc = len ( rgb ) <EOL> for i in range ( nc ) : <EOL> ctf . add_rgb_point ( rgb [ i ] [ <NUM_LIT:0> ] , * ( rgb [ i ] [ <NUM_LIT:1> : ] ) ) <EOL> if new_ctf : <EOL> volume_property . set_color ( ctf ) <EOL> try : <EOL> ctf . range = saved_data [ '<STR_LIT>' ] <EOL> except Exception : <EOL> pass <EOL> na = len ( a ) <EOL> new_otf = True <EOL> otf = volume_property . get_scalar_opacity ( ) <EOL> if isinstance ( otf , PiecewiseFunction ) : <EOL> new_otf = False <EOL> otf . remove_all_points ( ) <EOL> else : <EOL> otf = PiecewiseFunction ( ) <EOL> for i in range ( na ) : <EOL> otf . add_point ( a [ i ] [ <NUM_LIT:0> ] , a [ i ] [ <NUM_LIT:1> ] ) <EOL> if new_otf : <EOL> volume_property . set_scalar_opacity ( otf ) <EOL> return ctf , otf <EOL> def rescale_ctfs ( volume_property , new_range ) : <EOL> """<STR_LIT>""" <EOL> ctf = volume_property . rgb_transfer_function <EOL> otf = volume_property . get_scalar_opacity ( ) <EOL> old_range = ctf . range <EOL> def _rescale_value ( x , old , new ) : <EOL> nx = ( x - old [ <NUM_LIT:0> ] ) / ( old [ <NUM_LIT:1> ] - old [ <NUM_LIT:0> ] ) <EOL> return new [ <NUM_LIT:0> ] + nx * ( new [ <NUM_LIT:1> ] - new [ <NUM_LIT:0> ] ) <EOL> if new_range [ <NUM_LIT:0> ] != old_range [ <NUM_LIT:0> ] and new_range [ <NUM_LIT:1> ] != old_range [ <NUM_LIT:1> ] : <EOL> s_d = save_ctfs ( volume_property ) <EOL> s1 , s2 = new_range <EOL> if s1 > s2 : <EOL> s_d [ '<STR_LIT>' ] = ( s2 , s1 ) <EOL> else : <EOL> s_d [ '<STR_LIT>' ] = ( s1 , s2 ) <EOL> rgb = s_d [ '<STR_LIT>' ] <EOL> for v in rgb : <EOL> v [ <NUM_LIT:0> ] = _rescale_value ( v [ <NUM_LIT:0> ] , old_range , new_range ) <EOL> alpha = s_d [ '<STR_LIT>' ] <EOL> for v in alpha : <EOL> v [ <NUM_LIT:0> ] = _rescale_value ( v [ <NUM_LIT:0> ] , old_range , new_range ) <EOL> ctf , otf = load_ctfs ( s_d , volume_property ) <EOL> return ctf , otf <EOL> def set_lut ( lut , volume_property ) : <EOL> """<STR_LIT>""" <EOL> vp = volume_property <EOL> ctf = vp . rgb_transfer_function <EOL> otf = vp . get_scalar_opacity ( ) <EOL> s1 , s2 = ctf . range <EOL> nc = lut . number_of_colors <EOL> ds = float ( s2 - s1 ) / ( nc - <NUM_LIT:1> ) <EOL> for i in range ( nc ) : <EOL> r , g , b = ctf . get_color ( s1 + i * ds ) <EOL> a = otf . get_value ( s1 + i * ds ) <EOL> lut . set_table_value ( i , r , g , b , a ) <EOL> def set_ctf_from_lut ( lut , volume_property ) : <EOL> """<STR_LIT>""" <EOL> vp = volume_property <EOL> ctf = vp . rgb_transfer_function <EOL> s1 , s2 = ctf . range <EOL> nc = lut . number_of_colors <EOL> ds = float ( s2 - s1 ) / ( nc - <NUM_LIT:1> ) <EOL> ctf = ColorTransferFunction ( ) <EOL> otf = PiecewiseFunction ( ) <EOL> for i in range ( nc ) : <EOL> v = s1 + i * ds <EOL> r , g , b , a = lut . get_table_value ( i ) <EOL> ctf . add_rgb_point ( v , r , g , b ) <EOL> otf . add_point ( v , a ) <EOL> volume_property . set_color ( ctf ) <EOL> volume_property . set_scalar_opacity ( otf ) <EOL> class ColorTransferFunction ( tvtk . ColorTransferFunction ) : <EOL> """<STR_LIT>""" <EOL> nodes = List <EOL> def add_rgb_point ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> ret = super ( ColorTransferFunction , self ) . add_rgb_point ( * args ) <EOL> self . nodes . append ( args [ <NUM_LIT:0> ] ) <EOL> self . nodes . sort ( ) <EOL> return ret <EOL> def add_hsv_point ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> ret = super ( ColorTransferFunction , self ) . add_hsv_point ( * args ) <EOL> self . nodes . append ( args [ <NUM_LIT:0> ] ) <EOL> self . nodes . sort ( ) <EOL> return ret <EOL> def remove_all_points ( self ) : <EOL> """<STR_LIT>""" <EOL> super ( ColorTransferFunction , self ) . remove_all_points ( ) <EOL> self . nodes = [ ] <EOL> class PiecewiseFunction ( tvtk . PiecewiseFunction ) : <EOL> """<STR_LIT>""" <EOL> nodes = List <EOL> def initialize ( self ) : <EOL> """<STR_LIT>""" <EOL> super ( PiecewiseFunction , self ) . initialize ( ) <EOL> self . nodes = [ ] <EOL> def add_point ( self , x , val ) : <EOL> """<STR_LIT>""" <EOL> ret = super ( PiecewiseFunction , self ) . add_point ( x , val ) <EOL> self . nodes . append ( x ) <EOL> self . nodes . sort ( ) <EOL> return ret <EOL> def remove_point ( self , x ) : <EOL> """<STR_LIT>""" <EOL> ret = super ( PiecewiseFunction , self ) . remove_point ( x ) <EOL> self . nodes . remove ( x ) <EOL> self . nodes . sort ( ) <EOL> return ret <EOL> def remove_all_points ( self ) : <EOL> """<STR_LIT>""" <EOL> super ( PiecewiseFunction , self ) . remove_all_points ( ) <EOL> self . nodes = [ ] </s>
