Rendering osmesa (ofscreen mesa) opengl scene of ODE (physics engine) simulation

base.py

import abc

import six


@six.add_metaclass(abc.ABCMeta)
class Platform(object):
    """Base class for all OpenGL platforms.

    Parameters
    ----------
    viewport_width : int
        The width of the main viewport, in pixels.
    viewport_height : int
        The height of the main viewport, in pixels
    """

    def __init__(self, viewport_width, viewport_height):
        self.viewport_width = viewport_width
        self.viewport_height = viewport_height

    @property
    def viewport_width(self):
        """int : The width of the main viewport, in pixels.
        """
        return self._viewport_width

    @viewport_width.setter
    def viewport_width(self, value):
        self._viewport_width = value

    @property
    def viewport_height(self):
        """int : The height of the main viewport, in pixels.
        """
        return self._viewport_height

    @viewport_height.setter
    def viewport_height(self, value):
        self._viewport_height = value

    @abc.abstractmethod
    def init_context(self):
        """Create an OpenGL context.
        """
        pass

    @abc.abstractmethod
    def make_current(self):
        """Make the OpenGL context current.
        """
        pass

    @abc.abstractmethod
    def delete_context(self):
        """Delete the OpenGL context.
        """
        pass

    @abc.abstractmethod
    def supports_framebuffers(self):
        """Returns True if the method supports framebuffer rendering.
        """
        pass

    def __del__(self):
        try:
            self.delete_context()
        except Exception:
            pass

osmesa.py

from base import Platform


__all__ = ['OSMesaPlatform']


class OSMesaPlatform(Platform):
    """Renders into a software buffer using OSMesa. Requires special versions
    of OSMesa to be installed, plus PyOpenGL upgrade.
    """

    def __init__(self, viewport_width, viewport_height):
        super(OSMesaPlatform, self).__init__(viewport_width, viewport_height)
        self._context = None
        self._buffer = None

    def init_context(self, fixed_pipeline=False):
        '''

        :param fixed_pipeline: if enabled will set opengl 3.0, as it is the last version that supports fixed pipeline according to:
        https://www.khronos.org/opengl/wiki/Fixed_Function_Pipeline
        :return:
        '''
        from OpenGL import arrays
        from OpenGL.osmesa import (
            OSMesaCreateContextAttribs,
            OSMesaCreateContextExt,
            OSMESA_FORMAT,
            OSMESA_RGBA, OSMESA_PROFILE, OSMESA_CORE_PROFILE,
            OSMESA_CONTEXT_MAJOR_VERSION, OSMESA_CONTEXT_MINOR_VERSION,
            OSMESA_DEPTH_BITS
        )


        if fixed_pipeline:
            gl_version = (2,0)
        else:
            gl_version = (3, 3)

        attrs = arrays.GLintArray.asArray([
            OSMESA_FORMAT, OSMESA_RGBA,
            OSMESA_DEPTH_BITS, 24,
            OSMESA_PROFILE, OSMESA_CORE_PROFILE,
            OSMESA_CONTEXT_MAJOR_VERSION, gl_version[0],
            OSMESA_CONTEXT_MINOR_VERSION, gl_version[1],
            0
        ])
        ###self._context = OSMesaCreateContextAttribs(attrs, None)
        self._context = OSMesaCreateContextExt(OSMESA_RGBA, 24, 0, 0, None)
        self._buffer = arrays.GLubyteArray.zeros(
            (self.viewport_height, self.viewport_width, 4)
        )

    def make_current(self):
        from OpenGL import GL as gl
        from OpenGL.osmesa import OSMesaMakeCurrent
        assert(OSMesaMakeCurrent(
            self._context, self._buffer, gl.GL_UNSIGNED_BYTE,
            self.viewport_width, self.viewport_height
        ))

    def delete_context(self):
        from OpenGL.osmesa import OSMesaDestroyContext
        OSMesaDestroyContext(self._context)
        self._context = None
        self._buffer = None

    def supports_framebuffers(self):
        return False

tutorial3_osmesa.py

#!/usr/bin/env python

# http://pyode.sourceforge.net/tutorials/tutorial3.html

# pyODE example 3: Collision detection

# Originally by Matthias Baas.
# Updated by Pierre Gay to work without pygame or cgkit.
# Updated by Yoel Shoshan to work with osmesa (offscreen mesa) 

import sys, os, random, time
from math import *

os.environ['PYOPENGL_PLATFORM'] = 'osmesa'

from osmesa import OSMesaPlatform

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
import matplotlib.pyplot as plt
import numpy as np

import ode

plt.ion()

G_ax = None
#0.2,20)
Z_NEAR = 0.2     # Near clipping plane, in meters
Z_FAR = 20.0     # Far clipping plane, in meters

#last_znear = DEFAULT_Z_NEAR
#last_zfar = DEFAULT_Z_FAR

# geometric utility functions
def scalp (vec, scal):
    vec[0] *= scal
    vec[1] *= scal
    vec[2] *= scal

def length (vec):
    return sqrt (vec[0]**2 + vec[1]**2 + vec[2]**2)

# prepare_GL
def prepare_GL():
    """Prepare drawing.
    """

    # Viewport
    glViewport(0,0,640,480)

    # Initialize
    #glClearColor(0.8,0.8,0.9,0)
    glClearColor(0.0, 0.0, 0.0, 0)
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    glEnable(GL_DEPTH_TEST)
    glDisable(GL_LIGHTING)
    glEnable(GL_LIGHTING)
    glEnable(GL_NORMALIZE)
    glShadeModel(GL_FLAT)

    # Projection
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    gluPerspective(45,1.3333, Z_NEAR, Z_FAR)

    # Initialize ModelView matrix
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()

    # Light source
    glLightfv(GL_LIGHT0,GL_POSITION,[0,0,1,0])
    glLightfv(GL_LIGHT0,GL_DIFFUSE,[1,1,1,1])
    glLightfv(GL_LIGHT0,GL_SPECULAR,[1,1,1,1])
    glEnable(GL_LIGHT0)

    # View transformation
    gluLookAt (2.4, 3.6, 4.8, 0.5, 0.5, 0, 0, 1, 0)

# draw_body
def draw_body(body):
    """Draw an ODE body.
    """

    x,y,z = body.getPosition()
    R = body.getRotation()
    rot = [R[0], R[3], R[6], 0.,
           R[1], R[4], R[7], 0.,
           R[2], R[5], R[8], 0.,
           x, y, z, 1.0]
    glPushMatrix()
    glMultMatrixd(rot)
    if body.shape=="box":
        sx,sy,sz = body.boxsize
        glScalef(sx, sy, sz)
        glutSolidCube(1)
    glPopMatrix()


# create_box
def create_box(world, space, density, lx, ly, lz):
    """Create a box body and its corresponding geom."""

    # Create body
    body = ode.Body(world)
    M = ode.Mass()
    M.setBox(density, lx, ly, lz)
    body.setMass(M)

    # Set parameters for drawing the body
    body.shape = "box"
    body.boxsize = (lx, ly, lz)

    # Create a box geom for collision detection
    geom = ode.GeomBox(space, lengths=body.boxsize)
    geom.setBody(body)

    return body, geom

# drop_object
def drop_object():
    """Drop an object into the scene."""

    global bodies, geom, counter, objcount

    body, geom = create_box(world, space, 1000, 1.0,0.2,0.2)
    body.setPosition( (random.gauss(0,0.1),3.0,random.gauss(0,0.1)) )
    theta = random.uniform(0,2*pi)
    ct = cos (theta)
    st = sin (theta)
    body.setRotation([ct, 0., -st, 0., 1., 0., st, 0., ct])
    bodies.append(body)
    geoms.append(geom)
    counter=0
    objcount+=1

# explosion
def explosion():
    """Simulate an explosion.

    Every object is pushed away from the origin.
    The force is dependent on the objects distance from the origin.
    """
    global bodies

    for b in bodies:
        l=b.getPosition ()
        d = length (l)
        a = max(0, 40000*(1.0-0.2*d*d))
        l = [l[0] / 4, l[1], l[2] /4]
        scalp (l, a / length (l))
        b.addForce(l)

# pull
def pull():
    """Pull the objects back to the origin.

    Every object will be pulled back to the origin.
    Every couple of frames there'll be a thrust upwards so that
    the objects won't stick to the ground all the time.
    """
    global bodies, counter

    for b in bodies:
        l=list (b.getPosition ())
        scalp (l, -1000 / length (l))
        b.addForce(l)
        if counter%60==0:
            b.addForce((0,10000,0))

# Collision callback
def near_callback(args, geom1, geom2):
    """Callback function for the collide() method.

    This function checks if the given geoms do collide and
    creates contact joints if they do.
    """

    # Check if the objects do collide
    contacts = ode.collide(geom1, geom2)

    # Create contact joints
    world,contactgroup = args
    for c in contacts:
        c.setBounce(0.2)
        c.setMu(5000)
        j = ode.ContactJoint(world, contactgroup, c)
        j.attach(geom1.getBody(), geom2.getBody())



######################################################################

# Initialize Glut
glutInit ([])

# Open a window
#glutInitDisplayMode (GLUT_RGB | GLUT_DEPTH | GLUT_DOUBLE)

def read_color_buf(width, height):
    """Read and return the current viewport's color buffer.

    Alpha cannot be computed for an on-screen buffer.

    Returns
    -------
    color_im : (h, w, 3) uint8
        The color buffer in RGB byte format.
    """
    # Extract color image from frame buffer
    #width, height = self.viewport_width, self.viewport_height
    glBindFramebuffer(GL_READ_FRAMEBUFFER, 0)
    glReadBuffer(GL_FRONT)
    color_buf = glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE)

    # Re-format them into numpy arrays
    color_im = np.frombuffer(color_buf, dtype=np.uint8)
    color_im = color_im.reshape((height, width, 3))
    color_im = np.flip(color_im, axis=0)

    # Resize for macos if needed
    #if sys.platform == 'darwin':
    #    color_im = self._resize_image(color_im, True)

    return color_im

latest_zfar = None

def read_depth_buf(width, height):
    """Read and return the current viewport's color buffer.

    Returns
    -------

    depth_im : (h, w) float32
        The depth buffer in linear units.
    """
    

    #width, height = self.viewport_width, self.viewport_height
    glBindFramebuffer(GL_READ_FRAMEBUFFER, 0)
    glReadBuffer(GL_FRONT)
    depth_buf = glReadPixels(
        0, 0, width, height, GL_DEPTH_COMPONENT, GL_FLOAT
    )

    depth_im = np.frombuffer(depth_buf, dtype=np.float32)
    depth_im = depth_im.reshape((height, width))
    depth_im = np.flip(depth_im, axis=0)

    if False:

        inf_inds = (depth_im == 1.0)
        depth_im = 2.0 * depth_im - 1.0
        z_near, z_far = Z_NEAR, Z_FAR
        noninf = np.logical_not(inf_inds)
        if z_far is None:
            depth_im[noninf] = 2 * z_near / (1.0 - depth_im[noninf])
        else:
            depth_im[noninf] = ((2.0 * z_near * z_far) /
                                (z_far + z_near - depth_im[noninf] *
                                 (z_far - z_near)))
        depth_im[inf_inds] = 0.0

        # Resize for macos if needed
        #if sys.platform == 'darwin':
        #    depth_im = self._resize_image(depth_im)
    else:
        #depth_im -= Z_NEAR #depth_im.min()
        #depth_im /= Z_FAR #depth_im.max()
        #depth_im = (depth_im*255).astype(np.uint8)
        pass

    return depth_im

platform = OSMesaPlatform(640, 480)
platform.init_context(fixed_pipeline=True)
platform.make_current()


x = 0
y = 0
width = 640
height = 480
#glutInitWindowPosition (x, y);
#glutInitWindowSize (width, height);
#glutCreateWindow ("testode")

# Create a world object
world = ode.World()
world.setGravity( (0,-9.81,0) )
world.setERP(0.8)
world.setCFM(1E-5)

# Create a space object
space = ode.Space()

# Create a plane geom which prevent the objects from falling forever
floor = ode.GeomPlane(space, (0,1,0), 0)

# A list with ODE bodies
bodies = []

# The geoms for each of the bodies
geoms = []

# A joint group for the contact joints that are generated whenever
# two bodies collide
contactgroup = ode.JointGroup()

# Some variables used inside the simulation loop
fps = 50
dt = 1.0/fps
running = True
state = 0
counter = 0
objcount = 0
lasttime = time.time()


# keyboard callback
def _keyfunc (c, x, y):
    sys.exit (0)

#####glutKeyboardFunc (_keyfunc)

# draw callback
def _drawfunc ():
    global G_ax
    # Draw the scene
    prepare_GL()
    for b in bodies:
        draw_body(b)

    glFinish()

    color_buf = read_color_buf(width, height)

    #color_buf = read_depth_buf(width, height)
    #color_buf = np.vstack([color_buf[None, ...]] * 3).reshape(480, 640, 3)
    #color_buf = np.stack([color_buf] * 3, axis=-1)
    #color_buf = (color_buf*255.0).astype(np.uint8)
    #if G_ax is not None:
    #    G_ax.set_clim(vmin=color_buf.min(), vmax=color_buf.max())


    if G_ax is None:
        G_ax = plt.imshow(color_buf) #, 'gray')
    else:

        G_ax.set_data(color_buf)

    plt.pause(0.001)

    banana=123

##    glutSwapBuffers ()

##glutDisplayFunc (_drawfunc)

# idle callback
def _idlefunc ():
    global counter, state, lasttime

    t = dt - (time.time() - lasttime)
    if (t > 0):
        time.sleep(t)

    counter += 1

    if state==0:
        if counter==20:
            drop_object()
        if objcount==30:
            state=1
            counter=0
    # State 1: Explosion and pulling back the objects
    elif state==1:
        if counter==100:
            explosion()
        if counter>300:
            pull()
        if counter==500:
            counter=20

###    glutPostRedisplay ()

    # Simulate
    n = 4

    for i in range(n):
        # Detect collisions and create contact joints
        space.collide((world,contactgroup), near_callback)

        # Simulation step
        world.step(dt/n)

        # Remove all contact joints
        contactgroup.empty()

    lasttime = time.time()

##glutIdleFunc (_idlefunc)

###glutMainLoop ()

for i in range(2000):
    _idlefunc()
    _drawfunc()