import pyglet import pyglet.gl as gl import numpy as np import os import tempfile import subprocess import collections from trimesh import util from trimesh import transformations def previews(scene, resolution=(1080,1080), **kwargs): ''' Render a preview of a scene. Parameters ------------ scene: trimesh.Scene object resolution: (2,) int, resolution in pixels Returns --------- images: dict geometry name : bytes, PNG format ''' scene.set_camera() window = PreviewScene(scene,visible=True, resolution=resolution) for i in range(2): pyglet.clock.tick() window.switch_to() window.dispatch_events() window.dispatch_event('on_draw') window.flip() window.close() return window.renders def hex_to_rgb(color): value = str(color).lstrip('#').strip() if len(value) != 6: raise ValueError('hex colors must have 6 terms') rgb = [int(value[i:i+2], 16) for i in (0, 2 ,4)] return np.array(rgb) def camera_transform(centroid, extents, yaw=0.2, pitch=0.2): translation = np.eye(4) translation[0:3, 3] = centroid distance = ((extents.max() / 2) / np.tan(np.radians(60.0) / 2.0)) # offset by a distance set by the model size # the FOV is set for the Y axis, we multiply by a lightly # padded aspect ratio to make sure the model is in view initially translation[2][3] += distance * 1.35 transform = np.dot(transformations.rotation_matrix(yaw, [1, 0, 0], point=centroid), transformations.rotation_matrix(pitch, [0, 1, 0], point=centroid)) transform = np.linalg.inv(np.dot(transform, translation)) return transform class PreviewScene(pyglet.window.Window): def __init__(self, scene, visible=False, resolution=(640, 480)): self.scene = scene width, height = resolution conf = gl.Config(double_buffer=True) super(PreviewScene, self).__init__(config=conf, resizable=True, visible=visible, width=width, height=height) self.batch = pyglet.graphics.Batch() self.vertex_list = {} self.vertex_list_mode = {} for name, mesh in scene.geometry.items(): self.add_geometry(name=name, geometry=mesh) self.init_gl() self.set_size(*resolution) # what to render self.to_draw = collections.deque(scene.geometry.keys()) # make sure to render scene first self.to_draw.append('scene') self.renders = collections.OrderedDict() def _redraw(self): self.on_draw() def _add_mesh(self, name, mesh): self.vertex_list[name] = self.batch.add_indexed( *mesh_to_vertex_list(mesh)) self.vertex_list_mode[name] = gl.GL_TRIANGLES def _add_path(self, name, path): self.vertex_list[name] = self.batch.add_indexed( *path_to_vertex_list(path)) self.vertex_list_mode[name] = gl.GL_LINES def _add_points(self, name, pointcloud): self.vertex_list[name] = self.batch.add_indexed( *points_to_vertex_list(pointcloud.vertices, pointcloud.vertices_color)) self.vertex_list_mode[name] = gl.GL_POINTS def add_geometry(self, name, geometry): if util.is_instance_named(geometry, 'Trimesh'): return self._add_mesh(name, geometry) elif util.is_instance_named(geometry, 'Path3D'): return self._add_path(name, geometry) elif util.is_instance_named(geometry, 'Path2D'): return self._add_path(name, geometry.to_3D()) elif util.is_instance_named(geometry, 'PointCloud'): return self._add_points(name, geometry) else: raise ValueError('Geometry passed is not a viewable type!') def init_gl(self): # set background to a clear color if alpha is working # if alpha isn't working (AKA docker containers) set it # to an obscure light-ish shade of orange gl.glClearColor(*background_float) gl.glEnable(gl.GL_DEPTH_TEST) gl.glEnable(gl.GL_CULL_FACE) gl.glEnable(gl.GL_LIGHTING) gl.glEnable(gl.GL_LIGHT0) gl.glEnable(gl.GL_LIGHT1) gl.glLightfv(gl.GL_LIGHT0, gl.GL_POSITION, _gl_vector(.5, .5, 1, 0)) gl.glLightfv(gl.GL_LIGHT0, gl.GL_SPECULAR, _gl_vector(.5, .5, 1, 1)) gl.glLightfv(gl.GL_LIGHT0, gl.GL_DIFFUSE, _gl_vector(1, 1, 1, 1)) gl.glLightfv(gl.GL_LIGHT1, gl.GL_POSITION, _gl_vector(1, 0, .5, 0)) gl.glLightfv(gl.GL_LIGHT1, gl.GL_DIFFUSE, _gl_vector(.5, .5, .5, 1)) gl.glLightfv(gl.GL_LIGHT1, gl.GL_SPECULAR, _gl_vector(1, 1, 1, 1)) 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, _gl_vector(0.192250, 0.192250, 0.192250)) gl.glMaterialfv(gl.GL_FRONT, gl.GL_DIFFUSE, _gl_vector(0.507540, 0.507540, 0.507540)) gl.glMaterialfv(gl.GL_FRONT, gl.GL_SPECULAR, _gl_vector(.5082730, .5082730, .5082730)) gl.glMaterialf(gl.GL_FRONT, gl.GL_SHININESS, .4 * 128.0) gl.glEnable(gl.GL_BLEND) gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE_MINUS_SRC_ALPHA) gl.glEnable(gl.GL_LINE_SMOOTH) gl.glHint(gl.GL_LINE_SMOOTH_HINT, gl.GL_NICEST) gl.glLineWidth(1.5) gl.glPointSize(4) def on_resize(self, width, height): gl.glViewport(0, 0, width, height) gl.glMatrixMode(gl.GL_PROJECTION) gl.glLoadIdentity() gl.gluPerspective(60., width / float(height), .01, self.scene.scale * 5.0) gl.glMatrixMode(gl.GL_MODELVIEW) def on_draw(self): gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT) gl.glLoadIdentity() # pull the new camera transform from the scene transform_camera, junk = self.scene.graph['camera'] # apply the camera transform to the matrix stack gl.glMultMatrixf(_gl_matrix(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) count_original = len(node_names) count = -1 while len(node_names) > 0: count += 1 current_node = node_names.popleft() transform, geometry_name = self.scene.graph[current_node] if geometry_name is None: continue mesh = self.scene.geometry[geometry_name] if (hasattr(mesh, 'visual') and mesh.visual.transparency): # put the current item onto the back of the queue if count < count_original: node_names.append(current_node) continue # add a new matrix to the model stack gl.glPushMatrix() # transform by the nodes transform gl.glMultMatrixf(_gl_matrix(transform)) # get the mode of the current geometry mode = self.vertex_list_mode[geometry_name] # 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() self.save_image(name='scene') def save_image(self, name): colorbuffer = pyglet.image.get_buffer_manager().get_color_buffer() # if we want to modify the file we have to delete it ourselves later with tempfile.TemporaryFile() as f: colorbuffer.save(file=f) f.seek(0) self.renders[name] = f.read() def mesh_to_vertex_list(mesh): ''' Convert a Trimesh object to arguments for an indexed vertex list constructor. ''' vertex_count = len(mesh.triangles) * 3 normals = np.tile(mesh.face_normals, (1, 3)).reshape(-1).tolist() vertices = mesh.triangles.reshape(-1).tolist() faces = np.arange(vertex_count).tolist() colors = np.tile(mesh.visual.face_colors, (1,3)).reshape((-1,4)) color_gl = _validate_colors(colors, vertex_count) args = (vertex_count, # number of vertices gl.GL_TRIANGLES, # mode None, # group faces, # indices ('v3f/static', vertices), ('n3f/static', normals), color_gl) return args def path_to_vertex_list(path, group=None): vertices = path.vertices lines = np.vstack([util.stack_lines(e.discrete(path.vertices)) for e in path.entities]) index = np.arange(len(lines)) args = (len(lines), # number of vertices gl.GL_LINES, # mode group, # group index.reshape(-1).tolist(), # indices ('v3f/static', lines.reshape(-1)), ('c3f/static', np.array([.5, .10, .20] * len(lines)))) return args def points_to_vertex_list(points, colors, group=None): points = np.asanyarray(points) if not util.is_shape(points, (-1, 3)): raise ValueError('Pointcloud must be (n,3)!') color_gl = _validate_colors(colors, len(points)) index = np.arange(len(points)) args = (len(points), # number of vertices gl.GL_POINTS, # mode group, # group index.reshape(-1), # indices ('v3f/static', points.reshape(-1)), color_gl) return args def _validate_colors(colors, count): ''' Given a list of colors (or None) return a GL- acceptable list of colors Parameters ------------ colors: (count, (3 or 4)) colors Returns --------- colors_type: str, color type colors_gl: list, count length ''' colors = np.asanyarray(colors) count = int(count) if util.is_shape(colors, (count, (3, 4))): # convert the numpy dtype code to an opengl one colors_dtype = {'f': 'f', 'i': 'B', 'u': 'B'}[colors.dtype.kind] # create the data type description string pyglet expects colors_type = 'c' + str(colors.shape[1]) + colors_dtype + '/static' # reshape the 2D array into a 1D one and then convert to a python list colors = colors.reshape(-1).tolist() else: # case where colors are wrong shape, use a default color colors = np.tile([.5, .10, .20], (count, 1)).reshape(-1).tolist() colors_type = 'c3f/static' return colors_type, colors def _gl_matrix(array): ''' Convert a sane numpy transformation matrix (row major, (4,4)) to an stupid GLfloat transformation matrix (column major, (16,)) ''' a = np.array(array).T.reshape(-1) return (gl.GLfloat * len(a))(*a) def _gl_vector(array, *args): ''' Convert an array and an optional set of args into a flat vector of GLfloat ''' array = np.array(array) if len(args) > 0: array = np.append(array, args) vector = (gl.GLfloat * len(array))(*array) return vector background_hex = '#f9ede5' background_float = np.append(hex_to_rgb(background_hex) / 255.0, 0.0).tolist() if __name__ == '__main__': import trimesh mesh = trimesh.load('/home/mikedh/trimesh/models/cycloidal.3DXML') scene = trimesh.scene.split_scene(mesh) scene.convert_units('inches', guess=True) # function which manually runs the event loop render = previews(scene) from PIL import Image rendered = Image.open(trimesh.util.wrap_as_stream(render['scene'])) rendered.show()