Find file
Fetching contributors…
Cannot retrieve contributors at this time
114 lines (101 sloc) 4.25 KB
import os
import sys
import logging
import Imath
import OpenEXR
import numpy as np
import pyopencl as cl
import pyopencl.tools as cl_tools
from OpenGL.GL import *
log = logging.getLogger(__file__)
class ProgramCache(object):
def __init__(self, ctx):
from PyQt4.QtCore import QFileSystemWatcher
self.ctx = ctx
self.cache = {}
self.watcher = QFileSystemWatcher()
self.watcher.fileChanged.connect(self._file_changed)
self.options = '-I ' + os.path.join(os.path.dirname(__file__), 'kernels')
def _file_changed(self, path):
path = str(path)
del self.cache[path]
self.watcher.removePath(path)
def get(self, path):
prog = self.cache.get(path, None)
if not prog:
src = open(path).read()
prog = cl.Program(self.ctx, src).build(self.options)
log.info('%s succesfully compiled' % path)
self.cache[path] = prog
self.watcher.addPath(path)
return prog
def _create_texture(width, height):
tex = glGenTextures(1)
glBindTexture(GL_TEXTURE_2D, tex)
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA32F, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
return tex
def _create_context():
if sys.platform == 'darwin':
props = cl_tools.get_gl_sharing_context_properties()
ctx = cl.Context(properties=props, devices=[])
else:
props = [(cl.context_properties.PLATFORM, cl.get_platforms()[0])]
props += cl_tools.get_gl_sharing_context_properties()
ctx = cl.Context(properties=props)
for i, dev in enumerate(ctx.devices):
log.info('device %d: %s, driver version %s' % (i, dev.version, dev.driver_version))
log.info('device %d: %s by %s' % (i, dev.name, dev.vendor))
return ctx
def _read_channel(f, c):
return np.fromstring(f.channel(c, Imath.PixelType(Imath.PixelType.FLOAT)),
np.float32)
def _load_exr(ctx, filename):
f = OpenEXR.InputFile(filename)
dw = f.header()['dataWindow']
shape = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1)
sz = shape[0] * shape[1]
buf = np.empty(4 * sz, np.float32)
buf[0::4] = _read_channel(f, 'R')
buf[1::4] = _read_channel(f, 'G')
buf[2::4] = _read_channel(f, 'B')
buf[3::4] = np.ones(sz, np.float32)
return cl.Image(ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR,
cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.FLOAT),
shape=shape, hostbuf=buf)
class Orange(object):
def __init__(self, width=512, height=512):
self.sample = 0
self.width = width
self.height = height
self.tex = [_create_texture(width, height) for i in range(2)]
self.ctx = _create_context()
self.cache = ProgramCache(self.ctx)
self.queue = cl.CommandQueue(self.ctx)
mf = cl.mem_flags
self.buf = [cl.GLTexture(self.ctx, mf.READ_WRITE, GL_TEXTURE_2D, 0, t, 2) for t in self.tex]
self.front = 0
ident = np.array([1,0,0,0, 0,1,0,0, 0,0,1,0, 0,0,0,1], np.float32)
self.cam_xform = cl.Buffer(self.ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=ident)
self.env = _load_exr(self.ctx, 'images/MeadowTrail.exr')
def __del__(self):
glDeleteTextures(self.tex)
def render(self):
self.sample += 1
front_buf = self.buf[self.front]
back_buf = self.buf[1 - self.front]
self.front = 1 - self.front
prog = self.cache.get('kernels/test.cl')
cl.enqueue_acquire_gl_objects(self.queue, self.buf)
global_size = (self.width, self.height)
args = (self.cam_xform, np.int32(self.sample), self.env, back_buf, front_buf)
prog.test(self.queue, global_size, None, *args)
cl.enqueue_release_gl_objects(self.queue, self.buf)
self.queue.finish()
return front_buf.gl_object
def set_camera_xform(self, xform):
self.sample = 0
cl.enqueue_write_buffer(self.queue, self.cam_xform, np.array(xform, np.float32))