""" windowed.py --------------- Provides a pyglet- based windowed viewer to preview Trimesh, Scene, PointCloud, and Path objects. Works on all major platforms: Windows, Linux, and OSX. """ import platform import collections import numpy as np import pyglet import pyglet.gl as gl from .trackball import Trackball from .. import util from .. import rendering from ..visual import to_rgba from ..transformations import translation_matrix pyglet.options['shadow_window'] = False # smooth only when fewer faces than this _SMOOTH_MAX_FACES = 100000 class SceneViewer(pyglet.window.Window): def __init__(self, scene, smooth=True, flags=None, visible=True, resolution=None, start_loop=True, callback=None, callback_period=None, caption=None, fixed=None, offset_lines=True, background=None, window_conf=None, profile=False, ** kwargs): """ Create a window that will display a trimesh.Scene object in an OpenGL context via pyglet. Parameters --------------- scene : trimesh.scene.Scene Scene with geometry and transforms smooth : bool If True try to smooth shade things flags : dict If passed apply keys to self.view: ['cull', 'wireframe', etc] visible : bool Display window or not resolution : (2,) int Initial resolution of window start_loop : bool Call pyglet.app.run() at the end of init callback : function A function which can be called periodically to update things in the scene callback_period : float How often to call the callback, in seconds fixed : None or iterable List of keys in scene.geometry to skip view transform on to keep fixed relative to camera offset_lines : bool If True, will offset lines slightly so if drawn coplanar with mesh geometry they will be visible background : None or (4,) uint8 Color for background window_conf : None, or gl.Config Passed to window init kwargs : dict Additional arguments to pass, including 'background' for to set background color """ self.scene = self._scene = scene self.callback = callback self.callback_period = callback_period self.scene._redraw = self._redraw self.offset_lines = bool(offset_lines) self.background = background # save initial camera transform self._initial_camera_transform = scene.camera_transform.copy() # a transform to offset lines slightly to avoid Z-fighting if self.offset_lines: self._line_offset = translation_matrix( [0, 0, scene.scale / 1000]) self.reset_view(flags=flags) self.batch = pyglet.graphics.Batch() self._smooth = smooth self._profile = bool(profile) if self._profile: from pyinstrument import Profiler self.Profiler = Profiler # store kwargs self.kwargs = kwargs # store a vertexlist for an axis marker self._axis = None # store a vertexlist for a grid display self._grid = None # store scene geometry as vertex lists self.vertex_list = {} # store geometry hashes self.vertex_list_hash = {} # store geometry rendering mode self.vertex_list_mode = {} # store meshes that don't rotate relative to viewer self.fixed = fixed # store a hidden (don't not display) node. self._nodes_hidden = set() # name : texture self.textures = {} # if resolution isn't defined set a default value if resolution is None: resolution = scene.camera.resolution else: scene.camera.resolution = resolution # no window conf was passed so try to get the best looking one if window_conf is None: try: # try enabling antialiasing # if you have a graphics card this will probably work conf = gl.Config(sample_buffers=1, samples=4, depth_size=24, double_buffer=True) super(SceneViewer, self).__init__(config=conf, visible=visible, resizable=True, width=resolution[0], height=resolution[1], caption=caption) except pyglet.window.NoSuchConfigException: conf = gl.Config(double_buffer=True) super(SceneViewer, self).__init__(config=conf, resizable=True, visible=visible, width=resolution[0], height=resolution[1], caption=caption) else: # window config was manually passed super(SceneViewer, self).__init__(config=window_conf, resizable=True, visible=visible, width=resolution[0], height=resolution[1], caption=caption) # add scene geometry to viewer geometry self._update_vertex_list() # call after geometry is added self.init_gl() self.set_size(*resolution) self.update_flags() # someone has passed a callback to be called periodically if self.callback is not None: # if no callback period is specified set it to default if callback_period is None: # 30 times per second callback_period = 1.0 / 30.0 # set up a do-nothing periodic task which will # trigger `self.on_draw` every `callback_period` # seconds if someone has passed a callback pyglet.clock.schedule_interval(lambda x: x, callback_period) if start_loop: pyglet.app.run() def _redraw(self): self.on_draw() def _update_vertex_list(self): # update vertex_list if needed for name, geom in self.scene.geometry.items(): if geom.is_empty: continue if geometry_hash(geom) == self.vertex_list_hash.get(name): continue self.add_geometry(name=name, geometry=geom, smooth=bool(self._smooth)) def _update_meshes(self): # call the callback if specified if self.callback is not None: self.callback(self.scene) self._update_vertex_list() self._update_perspective(self.width, self.height) def add_geometry(self, name, geometry, **kwargs): """ Add a geometry to the viewer. Parameters -------------- name : hashable Name that references geometry geometry : Trimesh, Path2D, Path3D, PointCloud Geometry to display in the viewer window kwargs ** Passed to rendering.convert_to_vertexlist """ # convert geometry to constructor args args = rendering.convert_to_vertexlist(geometry, **kwargs) # create the indexed vertex list self.vertex_list[name] = self.batch.add_indexed(*args) # save the MD5 of the geometry self.vertex_list_hash[name] = geometry_hash(geometry) # save the rendering mode from the constructor args self.vertex_list_mode[name] = args[1] try: # if a geometry has UV coordinates that match vertices assert len(geometry.visual.uv) == len(geometry.vertices) has_tex = True except BaseException: has_tex = False if has_tex: tex = rendering.material_to_texture( geometry.visual.material) if tex is not None: self.textures[name] = tex def cleanup_geometries(self): """ Remove any stored vertex lists that no longer exist in the scene. """ # shorthand to scene graph graph = self.scene.graph # which parts of the graph still have geometry geom_keep = set([graph[node][1] for node in graph.nodes_geometry]) # which geometries no longer need to be kept geom_delete = [geom for geom in self.vertex_list if geom not in geom_keep] for geom in geom_delete: # remove stored vertex references self.vertex_list.pop(geom, None) self.vertex_list_hash.pop(geom, None) self.vertex_list_mode.pop(geom, None) self.textures.pop(geom, None) def unhide_geometry(self, node): """ If a node is hidden remove the flag and show the geometry on the next draw. Parameters ------------- node : str Node to display """ self._nodes_hidden.discard(node) def hide_geometry(self, node): """ Don't display the geometry contained at a node on the next draw. Parameters ------------- node : str Node to not display """ self._nodes_hidden.add(node) def reset_view(self, flags=None): """ Set view to the default view. Parameters -------------- flags : None or dict If any view key passed override the default e.g. {'cull': False} """ self.view = { 'cull': True, 'axis': False, 'grid': False, 'fullscreen': False, 'wireframe': False, 'ball': Trackball( pose=self._initial_camera_transform, size=self.scene.camera.resolution, scale=self.scene.scale, target=self.scene.centroid)} try: # if any flags are passed override defaults if isinstance(flags, dict): for k, v in flags.items(): if k in self.view: self.view[k] = v self.update_flags() except BaseException: pass def init_gl(self): """ Perform the magic incantations to create an OpenGL scene using pyglet. """ # if user passed a background color use it if self.background is None: # default background color is white background = np.ones(4) else: # convert to (4,) uint8 RGBA background = to_rgba(self.background) # convert to 0.0-1.0 float background = background.astype(np.float64) / 255.0 self._gl_set_background(background) # use camera setting for depth self._gl_enable_depth(self.scene.camera) self._gl_enable_color_material() self._gl_enable_blending() self._gl_enable_smooth_lines() self._gl_enable_lighting(self.scene) @staticmethod def _gl_set_background(background): gl.glClearColor(*background) @staticmethod def _gl_unset_background(): gl.glClearColor(*[0, 0, 0, 0]) @staticmethod def _gl_enable_depth(camera): """ Enable depth test in OpenGL using distances from `scene.camera`. """ # set the culling depth from our camera object gl.glDepthRange(camera.z_near, camera.z_far) gl.glClearDepth(1.0) gl.glEnable(gl.GL_DEPTH_TEST) gl.glDepthFunc(gl.GL_LEQUAL) gl.glEnable(gl.GL_DEPTH_TEST) gl.glEnable(gl.GL_CULL_FACE) @staticmethod def _gl_enable_color_material(): # do some openGL things gl.glColorMaterial(gl.GL_FRONT_AND_BACK, gl.GL_AMBIENT_AND_DIFFUSE) gl.glEnable(gl.GL_COLOR_MATERIAL) gl.glShadeModel(gl.GL_SMOOTH) gl.glMaterialfv(gl.GL_FRONT, gl.GL_AMBIENT, rendering.vector_to_gl( 0.192250, 0.192250, 0.192250)) gl.glMaterialfv(gl.GL_FRONT, gl.GL_DIFFUSE, rendering.vector_to_gl( 0.507540, 0.507540, 0.507540)) gl.glMaterialfv(gl.GL_FRONT, gl.GL_SPECULAR, rendering.vector_to_gl( .5082730, .5082730, .5082730)) gl.glMaterialf(gl.GL_FRONT, gl.GL_SHININESS, .4 * 128.0) @staticmethod def _gl_enable_blending(): # enable blending for transparency gl.glEnable(gl.GL_BLEND) gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE_MINUS_SRC_ALPHA) @staticmethod def _gl_enable_smooth_lines(): # make the lines from Path3D objects less ugly gl.glEnable(gl.GL_LINE_SMOOTH) gl.glHint(gl.GL_LINE_SMOOTH_HINT, gl.GL_NICEST) # set the width of lines to 4 pixels gl.glLineWidth(4) # set PointCloud markers to 4 pixels in size gl.glPointSize(4) @staticmethod def _gl_enable_lighting(scene): """ Take the lights defined in scene.lights and apply them as openGL lights. """ gl.glEnable(gl.GL_LIGHTING) # opengl only supports 7 lights? for i, light in enumerate(scene.lights[:7]): # the index of which light we have lightN = eval('gl.GL_LIGHT{}'.format(i)) # get the transform for the light by name matrix = scene.graph.get(light.name)[0] # convert light object to glLightfv calls multiargs = rendering.light_to_gl( light=light, transform=matrix, lightN=lightN) # enable the light in question gl.glEnable(lightN) # run the glLightfv calls for args in multiargs: gl.glLightfv(*args) def toggle_culling(self): """ Toggle back face culling. It is on by default but if you are dealing with non- watertight meshes you probably want to be able to see the back sides. """ self.view['cull'] = not self.view['cull'] self.update_flags() def toggle_wireframe(self): """ Toggle wireframe mode Good for looking inside meshes, off by default. """ self.view['wireframe'] = not self.view['wireframe'] self.update_flags() def toggle_fullscreen(self): """ Toggle between fullscreen and windowed mode. """ self.view['fullscreen'] = not self.view['fullscreen'] self.update_flags() def toggle_axis(self): """ Toggle a rendered XYZ/RGB axis marker: off, world frame, every frame """ # cycle through three axis states states = [False, 'world', 'all', 'without_world'] # the state after toggling index = (states.index(self.view['axis']) + 1) % len(states) # update state to next index self.view['axis'] = states[index] # perform gl actions self.update_flags() def toggle_grid(self): """ Toggle a rendered grid. """ # update state to next index self.view['grid'] = not self.view['grid'] # perform gl actions self.update_flags() def update_flags(self): """ Check the view flags, and call required GL functions. """ # view mode, filled vs wirefrom if self.view['wireframe']: gl.glPolygonMode(gl.GL_FRONT_AND_BACK, gl.GL_LINE) else: gl.glPolygonMode(gl.GL_FRONT_AND_BACK, gl.GL_FILL) # set fullscreen or windowed self.set_fullscreen(fullscreen=self.view['fullscreen']) # backface culling on or off if self.view['cull']: gl.glEnable(gl.GL_CULL_FACE) else: gl.glDisable(gl.GL_CULL_FACE) # case where we WANT an axis and NO vertexlist # is stored internally if self.view['axis'] and self._axis is None: from .. import creation # create an axis marker sized relative to the scene axis = creation.axis(origin_size=self.scene.scale / 100) # create ordered args for a vertex list args = rendering.mesh_to_vertexlist(axis) # store the axis as a reference self._axis = self.batch.add_indexed(*args) # case where we DON'T want an axis but a vertexlist # IS stored internally elif not self.view['axis'] and self._axis is not None: # remove the axis from the rendering batch self._axis.delete() # set the reference to None self._axis = None if self.view['grid'] and self._grid is None: try: # create a grid marker from ..path.creation import grid bounds = self.scene.bounds center = bounds.mean(axis=0) # set the grid to the lowest Z position # also offset by the scale to avoid interference center[2] = bounds[0][2] - (bounds[:, 2].ptp() / 100) # choose the side length by maximum XY length side = bounds.ptp(axis=0)[:2].max() # create an axis marker sized relative to the scene grid_mesh = grid( side=side, count=4, transform=translation_matrix(center)) # convert the path to vertexlist args args = rendering.convert_to_vertexlist(grid_mesh) # create ordered args for a vertex list self._grid = self.batch.add_indexed(*args) except BaseException: util.log.warning( 'failed to create grid!', exc_info=True) elif not self.view['grid'] and self._grid is not None: self._grid.delete() self._grid = None def _update_perspective(self, width, height): try: # for high DPI screens viewport size # will be different then the passed size width, height = self.get_viewport_size() except BaseException: # older versions of pyglet may not have this pass # set the new viewport size gl.glViewport(0, 0, width, height) gl.glMatrixMode(gl.GL_PROJECTION) gl.glLoadIdentity() # get field of view and Z range from camera camera = self.scene.camera # set perspective from camera data gl.gluPerspective(camera.fov[1], width / float(height), camera.z_near, camera.z_far) gl.glMatrixMode(gl.GL_MODELVIEW) return width, height def on_resize(self, width, height): """ Handle resized windows. """ width, height = self._update_perspective(width, height) self.scene.camera.resolution = (width, height) self.view['ball'].resize(self.scene.camera.resolution) self.scene.camera_transform[...] = self.view['ball'].pose def on_mouse_press(self, x, y, buttons, modifiers): """ Set the start point of the drag. """ self.view['ball'].set_state(Trackball.STATE_ROTATE) if (buttons == pyglet.window.mouse.LEFT): ctrl = (modifiers & pyglet.window.key.MOD_CTRL) shift = (modifiers & pyglet.window.key.MOD_SHIFT) if (ctrl and shift): self.view['ball'].set_state(Trackball.STATE_ZOOM) elif shift: self.view['ball'].set_state(Trackball.STATE_ROLL) elif ctrl: self.view['ball'].set_state(Trackball.STATE_PAN) elif (buttons == pyglet.window.mouse.MIDDLE): self.view['ball'].set_state(Trackball.STATE_PAN) elif (buttons == pyglet.window.mouse.RIGHT): self.view['ball'].set_state(Trackball.STATE_ZOOM) self.view['ball'].down(np.array([x, y])) self.scene.camera_transform[...] = self.view['ball'].pose def on_mouse_drag(self, x, y, dx, dy, buttons, modifiers): """ Pan or rotate the view. """ self.view['ball'].drag(np.array([x, y])) self.scene.camera_transform[...] = self.view['ball'].pose def on_mouse_scroll(self, x, y, dx, dy): """ Zoom the view. """ self.view['ball'].scroll(dy) self.scene.camera_transform[...] = self.view['ball'].pose def on_key_press(self, symbol, modifiers): """ Call appropriate functions given key presses. """ magnitude = 10 if symbol == pyglet.window.key.W: self.toggle_wireframe() elif symbol == pyglet.window.key.Z: self.reset_view() elif symbol == pyglet.window.key.C: self.toggle_culling() elif symbol == pyglet.window.key.A: self.toggle_axis() elif symbol == pyglet.window.key.G: self.toggle_grid() elif symbol == pyglet.window.key.Q: self.on_close() elif symbol == pyglet.window.key.M: self.maximize() elif symbol == pyglet.window.key.F: self.toggle_fullscreen() if symbol in [ pyglet.window.key.LEFT, pyglet.window.key.RIGHT, pyglet.window.key.DOWN, pyglet.window.key.UP]: self.view['ball'].down([0, 0]) if symbol == pyglet.window.key.LEFT: self.view['ball'].drag([-magnitude, 0]) elif symbol == pyglet.window.key.RIGHT: self.view['ball'].drag([magnitude, 0]) elif symbol == pyglet.window.key.DOWN: self.view['ball'].drag([0, -magnitude]) elif symbol == pyglet.window.key.UP: self.view['ball'].drag([0, magnitude]) self.scene.camera_transform[...] = self.view['ball'].pose def on_draw(self): """ Run the actual draw calls. """ if self._profile: profiler = self.Profiler() profiler.start() self._update_meshes() gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT) gl.glLoadIdentity() # pull the new camera transform from the scene transform_camera = np.linalg.inv(self.scene.camera_transform) # apply the camera transform to the matrix stack gl.glMultMatrixf(rendering.matrix_to_gl(transform_camera)) # we want to render fully opaque objects first, # followed by objects which have transparency node_names = collections.deque(self.scene.graph.nodes_geometry) # how many nodes did we start with count_original = len(node_names) count = -1 # if we are rendering an axis marker at the world if self._axis and not self.view['axis'] == 'without_world': # we stored it as a vertex list self._axis.draw(mode=gl.GL_TRIANGLES) if self._grid: self._grid.draw(mode=gl.GL_LINES) while len(node_names) > 0: count += 1 current_node = node_names.popleft() if current_node in self._nodes_hidden: continue # get the transform from world to geometry and mesh name transform, geometry_name = self.scene.graph.get(current_node) # if no geometry at this frame continue without rendering if geometry_name is None: continue # if a geometry is marked as fixed apply the inverse view transform if self.fixed is not None and geometry_name in self.fixed: # remove altered camera transform from fixed geometry transform_fix = np.linalg.inv( np.dot(self._initial_camera_transform, transform_camera)) # apply the transform so the fixed geometry doesn't move transform = np.dot(transform, transform_fix) # get a reference to the mesh so we can check transparency mesh = self.scene.geometry[geometry_name] if mesh.is_empty: continue # get the GL mode of the current geometry mode = self.vertex_list_mode[geometry_name] # if you draw a coplanar line with a triangle it will z-fight # the best way to do this is probably a shader but this works fine if mode == gl.GL_LINES: # apply the offset in camera space transform = util.multi_dot([ transform, np.linalg.inv(transform_camera), self._line_offset, transform_camera]) # add a new matrix to the model stack gl.glPushMatrix() # transform by the nodes transform gl.glMultMatrixf(rendering.matrix_to_gl(transform)) # draw an axis marker for each mesh frame if self.view['axis'] == 'all': self._axis.draw(mode=gl.GL_TRIANGLES) elif self.view['axis'] == 'without_world': if count > 0: self._axis.draw(mode=gl.GL_TRIANGLES) # transparent things must be drawn last if (hasattr(mesh, 'visual') and hasattr(mesh.visual, 'transparency') and mesh.visual.transparency): # put the current item onto the back of the queue if count < count_original: # add the node to be drawn last node_names.append(current_node) # pop the matrix stack for now gl.glPopMatrix() # come back to this mesh later continue # if we have texture enable the target texture texture = None if geometry_name in self.textures: texture = self.textures[geometry_name] gl.glEnable(texture.target) gl.glBindTexture(texture.target, texture.id) # draw the mesh with its transform applied self.vertex_list[geometry_name].draw(mode=mode) # pop the matrix stack as we drew what we needed to draw gl.glPopMatrix() # disable texture after using if texture is not None: gl.glDisable(texture.target) if self._profile: profiler.stop() print(profiler.output_text(unicode=True, color=True)) def save_image(self, file_obj): """ Save the current color buffer to a file object in PNG format. Parameters ------------- file_obj: file name, or file- like object """ manager = pyglet.image.get_buffer_manager() colorbuffer = manager.get_color_buffer() # if passed a string save by name if hasattr(file_obj, 'write'): colorbuffer.save(file=file_obj) else: colorbuffer.save(filename=file_obj) def geometry_hash(geometry): """ Get an MD5 for a geometry object Parameters ------------ geometry : object Returns ------------ MD5 : str """ if hasattr(geometry, 'md5'): # for most of our trimesh objects md5 = geometry.md5() elif hasattr(geometry, 'tostring'): # for unwrapped ndarray objects md5 = str(hash(geometry.tostring())) if hasattr(geometry, 'visual'): # if visual properties are defined md5 += str(geometry.visual.crc()) return md5 def render_scene(scene, resolution=None, visible=True, **kwargs): """ Render a preview of a scene to a PNG. Parameters ------------ scene : trimesh.Scene Geometry to be rendered resolution : (2,) int or None Resolution in pixels, or set from scene.camera kwargs : ** Passed to SceneViewer Returns --------- render : bytes Image in PNG format """ window = SceneViewer(scene, start_loop=False, visible=visible, resolution=resolution, **kwargs) if visible is None: visible = platform.system() != 'Linux' from ..util import BytesIO # need to run loop twice to display anything for save in [False, False, True]: pyglet.clock.tick() window.switch_to() window.dispatch_events() window.dispatch_event('on_draw') window.flip() if save: # save the color buffer data to memory file_obj = BytesIO() window.save_image(file_obj) file_obj.seek(0) render = file_obj.read() window.close() return render