# -*- coding: utf-8 -*- """Mayavi/traits GUI visualization elements.""" # Authors: Christian Brodbeck # # License: BSD (3-clause) import numpy as np from mayavi.mlab import pipeline, text3d from mayavi.modules.glyph import Glyph from mayavi.modules.surface import Surface from mayavi.sources.vtk_data_source import VTKDataSource from mayavi.tools.mlab_scene_model import MlabSceneModel from traits.api import (HasTraits, HasPrivateTraits, on_trait_change, Instance, Array, Bool, Button, Enum, Float, Int, List, Range, Str, RGBColor, Property, cached_property) from traitsui.api import (View, Item, HGroup, VGrid, VGroup, Spring, TextEditor) from tvtk.api import tvtk from ..defaults import DEFAULTS from ..surface import _project_onto_surface, _normalize_vectors from ..source_space import _points_outside_surface from ..utils import SilenceStdout from ..viz._3d import _create_mesh_surf, _toggle_mlab_render headview_borders = VGroup(Item('headview', style='custom', show_label=False), show_border=True, label='View') def _mm_fmt(x): """Format data in units of mm.""" return '%0.1f' % x laggy_float_editor_mm = TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x: '%0.1f' % x) laggy_float_editor_scale = TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x: '%0.1f' % x) laggy_float_editor_headscale = TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x: '%0.3f' % x) laggy_float_editor_weight = TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x: '%0.2f' % x) laggy_float_editor_scale = TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x: '%0.1f' % x) laggy_float_editor_deg = TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x: '%0.1f' % x) _BUTTON_WIDTH = -80 _DEG_WIDTH = -50 # radian floats _MM_WIDTH = _DEG_WIDTH # mm floats _SCALE_WIDTH = _DEG_WIDTH # scale floats _INC_BUTTON_WIDTH = -25 # inc/dec buttons _DEG_STEP_WIDTH = -50 _MM_STEP_WIDTH = _DEG_STEP_WIDTH _SCALE_STEP_WIDTH = _DEG_STEP_WIDTH _WEIGHT_WIDTH = -60 # weight floats _VIEW_BUTTON_WIDTH = -60 # width is optimized for macOS and Linux avoid a horizontal scroll-bar # even when a vertical one is present _COREG_WIDTH = -290 _TEXT_WIDTH = -260 _REDUCED_TEXT_WIDTH = _TEXT_WIDTH - 40 * np.sign(_TEXT_WIDTH) _DIG_SOURCE_WIDTH = _TEXT_WIDTH - 50 * np.sign(_TEXT_WIDTH) _MRI_FIDUCIALS_WIDTH = _TEXT_WIDTH - 60 * np.sign(_TEXT_WIDTH) _SHOW_BORDER = True _RESET_LABEL = u"↻" _RESET_WIDTH = _INC_BUTTON_WIDTH class HeadViewController(HasTraits): """Set head views for the given coordinate system. Parameters ---------- system : 'RAS' | 'ALS' | 'ARI' Coordinate system described as initials for directions associated with the x, y, and z axes. Relevant terms are: Anterior, Right, Left, Superior, Inferior. """ system = Enum("RAS", "ALS", "ARI", desc="Coordinate system: directions of " "the x, y, and z axis.") right = Button() front = Button() left = Button() top = Button() interaction = Enum('trackball', 'terrain') scale = Float(0.16) scene = Instance(MlabSceneModel) view = View(VGroup( VGrid('0', Item('top', width=_VIEW_BUTTON_WIDTH), '0', Item('right', width=_VIEW_BUTTON_WIDTH), Item('front', width=_VIEW_BUTTON_WIDTH), Item('left', width=_VIEW_BUTTON_WIDTH), columns=3, show_labels=False), '_', HGroup(Item('scale', label='Scale', editor=laggy_float_editor_headscale, width=_SCALE_WIDTH, show_label=True), Item('interaction', tooltip='Mouse interaction mode', show_label=False), Spring()), show_labels=False)) @on_trait_change('scene.activated') def _init_view(self): self.scene.parallel_projection = True # apparently scene,activated happens several times if self.scene.renderer: self.sync_trait('scale', self.scene.camera, 'parallel_scale') # and apparently this does not happen by default: self.on_trait_change(self.scene.render, 'scale') self.interaction = self.interaction # could be delayed @on_trait_change('interaction') def on_set_interaction(self, _, interaction): if self.scene is None or self.scene.interactor is None: return # Ensure we're in the correct orientation for the # InteractorStyleTerrain to have the correct "up" self.on_set_view('front', '') self.scene.mlab.draw() self.scene.interactor.interactor_style = \ tvtk.InteractorStyleTerrain() if interaction == 'terrain' else \ tvtk.InteractorStyleTrackballCamera() # self.scene.interactor.interactor_style. self.on_set_view('front', '') self.scene.mlab.draw() @on_trait_change('top,left,right,front') def on_set_view(self, view, _): if self.scene is None: return system = self.system kwargs = dict(ALS=dict(front=(0, 90, -90), left=(90, 90, 180), right=(-90, 90, 0), top=(0, 0, -90)), RAS=dict(front=(90., 90., 180), left=(180, 90, 90), right=(0., 90, 270), top=(90, 0, 180)), ARI=dict(front=(0, 90, 90), left=(-90, 90, 180), right=(90, 90, 0), top=(0, 180, 90))) if system not in kwargs: raise ValueError("Invalid system: %r" % system) if view not in kwargs[system]: raise ValueError("Invalid view: %r" % view) kwargs = dict(zip(('azimuth', 'elevation', 'roll'), kwargs[system][view])) kwargs['focalpoint'] = (0., 0., 0.) with SilenceStdout(): self.scene.mlab.view(distance=None, reset_roll=True, figure=self.scene.mayavi_scene, **kwargs) class Object(HasPrivateTraits): """Represent a 3d object in a mayavi scene.""" points = Array(float, shape=(None, 3)) nn = Array(float, shape=(None, 3)) name = Str scene = Instance(MlabSceneModel, ()) src = Instance(VTKDataSource) # This should be Tuple, but it is broken on Anaconda as of 2016/12/16 color = RGBColor((1., 1., 1.)) # Due to a MESA bug, we use 0.99 opacity to force alpha blending opacity = Range(low=0., high=1., value=0.99) visible = Bool(True) def _update_points(self): """Update the location of the plotted points.""" if hasattr(self.src, 'data'): self.src.data.points = self.points return True class PointObject(Object): """Represent a group of individual points in a mayavi scene.""" label = Bool(False) label_scale = Float(0.01) projectable = Bool(False) # set based on type of points orientable = Property(depends_on=['project_to_points']) text3d = List point_scale = Float(10, label='Point Scale') # projection onto a surface project_to_points = Array(float, shape=(None, 3)) project_to_tris = Array(int, shape=(None, 3)) project_to_surface = Bool(False, label='Project', desc='project points ' 'onto the surface') orient_to_surface = Bool(False, label='Orient', desc='orient points ' 'toward the surface') scale_by_distance = Bool(False, label='Dist.', desc='scale points by ' 'distance from the surface') mark_inside = Bool(False, label='Mark', desc='mark points inside the ' 'surface in a different color') inside_color = RGBColor((0., 0., 0.)) glyph = Instance(Glyph) resolution = Int(8) view = View(HGroup(Item('visible', show_label=False), Item('color', show_label=False), Item('opacity'))) def __init__(self, view='points', has_norm=False, *args, **kwargs): """Init. Parameters ---------- view : 'points' | 'cloud' Whether the view options should be tailored to individual points or a point cloud. has_norm : bool Whether a norm can be defined; adds view options based on point norms (default False). """ assert view in ('points', 'cloud', 'arrow') self._view = view self._has_norm = bool(has_norm) super(PointObject, self).__init__(*args, **kwargs) def default_traits_view(self): # noqa: D102 color = Item('color', show_label=False) scale = Item('point_scale', label='Size', width=_SCALE_WIDTH, editor=laggy_float_editor_headscale) orient = Item('orient_to_surface', enabled_when='orientable and not project_to_surface', tooltip='Orient points toward the surface') dist = Item('scale_by_distance', enabled_when='orientable and not project_to_surface', tooltip='Scale points by distance from the surface') mark = Item('mark_inside', enabled_when='orientable and not project_to_surface', tooltip='Mark points inside the surface using a different ' 'color') if self._view == 'arrow': visible = Item('visible', label='Show', show_label=False) return View(HGroup(visible, scale, 'opacity', 'label', Spring())) elif self._view == 'points': visible = Item('visible', label='Show', show_label=True) views = (visible, color, scale, 'label') else: assert self._view == 'cloud' visible = Item('visible', show_label=False) views = (visible, color, scale) if not self._has_norm: return View(HGroup(*views)) group2 = HGroup(dist, Item('project_to_surface', show_label=True, enabled_when='projectable', tooltip='Project points onto the surface ' '(for visualization, does not affect ' 'fitting)'), orient, mark, Spring(), show_left=False) return View(HGroup(HGroup(*views), group2)) @on_trait_change('label') def _show_labels(self, show): _toggle_mlab_render(self, False) while self.text3d: text = self.text3d.pop() text.remove() if show and len(self.src.data.points) > 0: fig = self.scene.mayavi_scene if self._view == 'arrow': # for axes x, y, z = self.src.data.points[0] self.text3d.append(text3d( x, y, z, self.name, scale=self.label_scale, color=self.color, figure=fig)) else: for i, (x, y, z) in enumerate(np.array(self.src.data.points)): self.text3d.append(text3d( x, y, z, ' %i' % i, scale=self.label_scale, color=self.color, figure=fig)) _toggle_mlab_render(self, True) @on_trait_change('visible') def _on_hide(self): if not self.visible: self.label = False @on_trait_change('scene.activated') def _plot_points(self): """Add the points to the mayavi pipeline""" if self.scene is None: return if hasattr(self.glyph, 'remove'): self.glyph.remove() if hasattr(self.src, 'remove'): self.src.remove() _toggle_mlab_render(self, False) x, y, z = self.points.T fig = self.scene.mayavi_scene scatter = pipeline.scalar_scatter(x, y, z, fig=fig) if not scatter.running: # this can occur sometimes during testing w/ui.dispose() return # fig.scene.engine.current_object is scatter mode = 'arrow' if self._view == 'arrow' else 'sphere' glyph = pipeline.glyph(scatter, color=self.color, figure=fig, scale_factor=self.point_scale, opacity=1., resolution=self.resolution, mode=mode) glyph.actor.property.backface_culling = True glyph.glyph.glyph.vector_mode = 'use_normal' glyph.glyph.glyph.clamping = False if mode == 'arrow': glyph.glyph.glyph_source.glyph_position = 'tail' glyph.actor.mapper.color_mode = 'map_scalars' glyph.actor.mapper.scalar_mode = 'use_point_data' glyph.actor.mapper.use_lookup_table_scalar_range = False self.src = scatter self.glyph = glyph self.sync_trait('point_scale', self.glyph.glyph.glyph, 'scale_factor') self.sync_trait('color', self.glyph.actor.property, mutual=False) self.sync_trait('visible', self.glyph) self.sync_trait('opacity', self.glyph.actor.property) self.sync_trait('mark_inside', self.glyph.actor.mapper, 'scalar_visibility') self.on_trait_change(self._update_points, 'points') self._update_marker_scaling() self._update_marker_type() self._update_colors() _toggle_mlab_render(self, True) # self.scene.camera.parallel_scale = _scale # don't put project_to_tris here, just always set project_to_points second @on_trait_change('points,project_to_points,project_to_surface,mark_inside') def _update_projections(self): """Update the styles of the plotted points.""" if not hasattr(self.src, 'data'): return if self._view == 'arrow': self.src.data.point_data.normals = self.nn self.src.data.point_data.update() return # projections if len(self.project_to_points) <= 1 or len(self.points) == 0: return # Do the projections pts = self.points surf = dict(rr=np.array(self.project_to_points), tris=np.array(self.project_to_tris)) method = 'accurate' if len(surf['rr']) <= 20484 else 'nearest' proj_pts, proj_nn = _project_onto_surface( pts, surf, project_rrs=True, return_nn=True, method=method)[2:4] vec = pts - proj_pts # point to the surface if self.project_to_surface: pts = proj_pts nn = proj_nn else: nn = vec.copy() _normalize_vectors(nn) if self.mark_inside and not self.project_to_surface: scalars = _points_outside_surface(pts, surf).astype(int) else: scalars = np.ones(len(pts)) # With this, a point exactly on the surface is of size point_scale dist = np.linalg.norm(vec, axis=-1, keepdims=True) self.src.data.point_data.normals = (250 * dist + 1) * nn self.src.data.point_data.scalars = scalars self.glyph.actor.mapper.scalar_range = [0., 1.] self.src.data.points = pts # projection can change this self.src.data.point_data.update() @on_trait_change('color,inside_color') def _update_colors(self): if self.glyph is None: return # inside_color is the surface color, let's try to get far # from that inside = np.array(self.inside_color) # if it's too close to gray, just use black: if np.mean(np.abs(inside - 0.5)) < 0.2: inside.fill(0.) else: inside = 1 - inside colors = np.array([tuple(inside) + (1,), tuple(self.color) + (1,)]) * 255. self.glyph.module_manager.scalar_lut_manager.lut.table = colors @on_trait_change('project_to_surface,orient_to_surface') def _update_marker_type(self): # not implemented for arrow if self.glyph is None or self._view == 'arrow': return defaults = DEFAULTS['coreg'] gs = self.glyph.glyph.glyph_source res = getattr(gs.glyph_source, 'theta_resolution', getattr(gs.glyph_source, 'resolution', None)) if self.project_to_surface or self.orient_to_surface: gs.glyph_source = tvtk.CylinderSource() gs.glyph_source.height = defaults['eegp_height'] gs.glyph_source.center = (0., -defaults['eegp_height'], 0) gs.glyph_source.resolution = res else: gs.glyph_source = tvtk.SphereSource() gs.glyph_source.phi_resolution = res gs.glyph_source.theta_resolution = res @on_trait_change('scale_by_distance,project_to_surface') def _update_marker_scaling(self): if self.glyph is None: return if self.scale_by_distance and not self.project_to_surface: self.glyph.glyph.scale_mode = 'scale_by_vector' else: self.glyph.glyph.scale_mode = 'data_scaling_off' def _resolution_changed(self, new): if not self.glyph: return gs = self.glyph.glyph.glyph_source.glyph_source if isinstance(gs, tvtk.SphereSource): gs.phi_resolution = new gs.theta_resolution = new elif isinstance(gs, tvtk.CylinderSource): gs.resolution = new else: # ArrowSource gs.tip_resolution = new gs.shaft_resolution = new @cached_property def _get_orientable(self): return (len(self.project_to_points) > 0 and len(self.project_to_tris) > 0) class SurfaceObject(Object): """Represent a solid object in a mayavi scene. Notes ----- Doesn't automatically update plot because update requires both :attr:`points` and :attr:`tris`. Call :meth:`plot` after updating both attributes. """ rep = Enum("Surface", "Wireframe") tris = Array(int, shape=(None, 3)) surf = Instance(Surface) surf_rear = Instance(Surface) view = View(HGroup(Item('visible', show_label=False), Item('color', show_label=False), Item('opacity'))) def __init__(self, block_behind=False, **kwargs): # noqa: D102 self._block_behind = block_behind self._deferred_tri_update = False super(SurfaceObject, self).__init__(**kwargs) def clear(self): # noqa: D102 if hasattr(self.src, 'remove'): self.src.remove() if hasattr(self.surf, 'remove'): self.surf.remove() if hasattr(self.surf_rear, 'remove'): self.surf_rear.remove() self.reset_traits(['src', 'surf']) @on_trait_change('scene.activated') def plot(self): """Add the points to the mayavi pipeline""" _scale = self.scene.camera.parallel_scale self.clear() if not np.any(self.tris): return fig = self.scene.mayavi_scene surf = dict(rr=self.points, tris=self.tris) normals = _create_mesh_surf(surf, fig=fig) self.src = normals.parent rep = 'wireframe' if self.rep == 'Wireframe' else 'surface' # Add the opaque "inside" first to avoid the translucent "outside" # from being occluded (gh-5152) if self._block_behind: surf_rear = pipeline.surface( normals, figure=fig, color=self.color, representation=rep, line_width=1) surf_rear.actor.property.frontface_culling = True self.surf_rear = surf_rear self.sync_trait('color', self.surf_rear.actor.property, mutual=False) self.sync_trait('visible', self.surf_rear, 'visible') self.surf_rear.actor.property.opacity = 1. surf = pipeline.surface( normals, figure=fig, color=self.color, representation=rep, line_width=1) surf.actor.property.backface_culling = True self.surf = surf self.sync_trait('visible', self.surf, 'visible') self.sync_trait('color', self.surf.actor.property, mutual=False) self.sync_trait('opacity', self.surf.actor.property) self.scene.camera.parallel_scale = _scale @on_trait_change('tris') def _update_tris(self): self._deferred_tris_update = True @on_trait_change('points') def _update_points(self): if Object._update_points(self): if self._deferred_tris_update: self.src.data.polys = self.tris self._deffered_tris_update = False self.src.update() # necessary for SurfaceObject since Mayavi 4.5.0