Manual pyglet event loop headless rendering

headless.py

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()