Spaces:
Running
on
T4
Running
on
T4
"""OpenGL shader program wrapper. | |
""" | |
import numpy as np | |
import os | |
import re | |
import OpenGL | |
from OpenGL.GL import * | |
from OpenGL.GL import shaders as gl_shader_utils | |
class ShaderProgramCache(object): | |
"""A cache for shader programs. | |
""" | |
def __init__(self, shader_dir=None): | |
self._program_cache = {} | |
self.shader_dir = shader_dir | |
if self.shader_dir is None: | |
base_dir, _ = os.path.split(os.path.realpath(__file__)) | |
self.shader_dir = os.path.join(base_dir, 'shaders') | |
def get_program(self, vertex_shader, fragment_shader, | |
geometry_shader=None, defines=None): | |
"""Get a program via a list of shader files to include in the program. | |
Parameters | |
---------- | |
vertex_shader : str | |
The vertex shader filename. | |
fragment_shader : str | |
The fragment shader filename. | |
geometry_shader : str | |
The geometry shader filename. | |
defines : dict | |
Defines and their values for the shader. | |
Returns | |
------- | |
program : :class:`.ShaderProgram` | |
The program. | |
""" | |
shader_names = [] | |
if defines is None: | |
defines = {} | |
shader_filenames = [ | |
x for x in [vertex_shader, fragment_shader, geometry_shader] | |
if x is not None | |
] | |
for fn in shader_filenames: | |
if fn is None: | |
continue | |
_, name = os.path.split(fn) | |
shader_names.append(name) | |
cid = OpenGL.contextdata.getContext() | |
key = tuple([cid] + sorted( | |
[(s,1) for s in shader_names] + [(d, defines[d]) for d in defines] | |
)) | |
if key not in self._program_cache: | |
shader_filenames = [ | |
os.path.join(self.shader_dir, fn) for fn in shader_filenames | |
] | |
if len(shader_filenames) == 2: | |
shader_filenames.append(None) | |
vs, fs, gs = shader_filenames | |
self._program_cache[key] = ShaderProgram( | |
vertex_shader=vs, fragment_shader=fs, | |
geometry_shader=gs, defines=defines | |
) | |
return self._program_cache[key] | |
def clear(self): | |
for key in self._program_cache: | |
self._program_cache[key].delete() | |
self._program_cache = {} | |
class ShaderProgram(object): | |
"""A thin wrapper about OpenGL shader programs that supports easy creation, | |
binding, and uniform-setting. | |
Parameters | |
---------- | |
vertex_shader : str | |
The vertex shader filename. | |
fragment_shader : str | |
The fragment shader filename. | |
geometry_shader : str | |
The geometry shader filename. | |
defines : dict | |
Defines and their values for the shader. | |
""" | |
def __init__(self, vertex_shader, fragment_shader, | |
geometry_shader=None, defines=None): | |
self.vertex_shader = vertex_shader | |
self.fragment_shader = fragment_shader | |
self.geometry_shader = geometry_shader | |
self.defines = defines | |
if self.defines is None: | |
self.defines = {} | |
self._program_id = None | |
self._vao_id = None # PYOPENGL BUG | |
# DEBUG | |
# self._unif_map = {} | |
def _add_to_context(self): | |
if self._program_id is not None: | |
raise ValueError('Shader program already in context') | |
shader_ids = [] | |
# Load vert shader | |
shader_ids.append(gl_shader_utils.compileShader( | |
self._load(self.vertex_shader), GL_VERTEX_SHADER) | |
) | |
# Load frag shader | |
shader_ids.append(gl_shader_utils.compileShader( | |
self._load(self.fragment_shader), GL_FRAGMENT_SHADER) | |
) | |
# Load geometry shader | |
if self.geometry_shader is not None: | |
shader_ids.append(gl_shader_utils.compileShader( | |
self._load(self.geometry_shader), GL_GEOMETRY_SHADER) | |
) | |
# Bind empty VAO PYOPENGL BUG | |
if self._vao_id is None: | |
self._vao_id = glGenVertexArrays(1) | |
glBindVertexArray(self._vao_id) | |
# Compile program | |
self._program_id = gl_shader_utils.compileProgram(*shader_ids) | |
# Unbind empty VAO PYOPENGL BUG | |
glBindVertexArray(0) | |
def _in_context(self): | |
return self._program_id is not None | |
def _remove_from_context(self): | |
if self._program_id is not None: | |
glDeleteProgram(self._program_id) | |
glDeleteVertexArrays(1, [self._vao_id]) | |
self._program_id = None | |
self._vao_id = None | |
def _load(self, shader_filename): | |
path, _ = os.path.split(shader_filename) | |
with open(shader_filename) as f: | |
text = f.read() | |
def ifdef(matchobj): | |
if matchobj.group(1) in self.defines: | |
return '#if 1' | |
else: | |
return '#if 0' | |
def ifndef(matchobj): | |
if matchobj.group(1) in self.defines: | |
return '#if 0' | |
else: | |
return '#if 1' | |
ifdef_regex = re.compile( | |
'#ifdef\\s+([a-zA-Z_][a-zA-Z_0-9]*)\\s*$', re.MULTILINE | |
) | |
ifndef_regex = re.compile( | |
'#ifndef\\s+([a-zA-Z_][a-zA-Z_0-9]*)\\s*$', re.MULTILINE | |
) | |
text = re.sub(ifdef_regex, ifdef, text) | |
text = re.sub(ifndef_regex, ifndef, text) | |
for define in self.defines: | |
value = str(self.defines[define]) | |
text = text.replace(define, value) | |
return text | |
def _bind(self): | |
"""Bind this shader program to the current OpenGL context. | |
""" | |
if self._program_id is None: | |
raise ValueError('Cannot bind program that is not in context') | |
# glBindVertexArray(self._vao_id) | |
glUseProgram(self._program_id) | |
def _unbind(self): | |
"""Unbind this shader program from the current OpenGL context. | |
""" | |
glUseProgram(0) | |
def delete(self): | |
"""Delete this shader program from the current OpenGL context. | |
""" | |
self._remove_from_context() | |
def set_uniform(self, name, value, unsigned=False): | |
"""Set a uniform value in the current shader program. | |
Parameters | |
---------- | |
name : str | |
Name of the uniform to set. | |
value : int, float, or ndarray | |
Value to set the uniform to. | |
unsigned : bool | |
If True, ints will be treated as unsigned values. | |
""" | |
try: | |
# DEBUG | |
# self._unif_map[name] = 1, (1,) | |
loc = glGetUniformLocation(self._program_id, name) | |
if loc == -1: | |
raise ValueError('Invalid shader variable: {}'.format(name)) | |
if isinstance(value, np.ndarray): | |
# DEBUG | |
# self._unif_map[name] = value.size, value.shape | |
if value.ndim == 1: | |
if (np.issubdtype(value.dtype, np.unsignedinteger) or | |
unsigned): | |
dtype = 'u' | |
value = value.astype(np.uint32) | |
elif np.issubdtype(value.dtype, np.integer): | |
dtype = 'i' | |
value = value.astype(np.int32) | |
else: | |
dtype = 'f' | |
value = value.astype(np.float32) | |
self._FUNC_MAP[(value.shape[0], dtype)](loc, 1, value) | |
else: | |
self._FUNC_MAP[(value.shape[0], value.shape[1])]( | |
loc, 1, GL_TRUE, value | |
) | |
# Call correct uniform function | |
elif isinstance(value, float): | |
glUniform1f(loc, value) | |
elif isinstance(value, int): | |
if unsigned: | |
glUniform1ui(loc, value) | |
else: | |
glUniform1i(loc, value) | |
elif isinstance(value, bool): | |
if unsigned: | |
glUniform1ui(loc, int(value)) | |
else: | |
glUniform1i(loc, int(value)) | |
else: | |
raise ValueError('Invalid data type') | |
except Exception: | |
pass | |
_FUNC_MAP = { | |
(1,'u'): glUniform1uiv, | |
(2,'u'): glUniform2uiv, | |
(3,'u'): glUniform3uiv, | |
(4,'u'): glUniform4uiv, | |
(1,'i'): glUniform1iv, | |
(2,'i'): glUniform2iv, | |
(3,'i'): glUniform3iv, | |
(4,'i'): glUniform4iv, | |
(1,'f'): glUniform1fv, | |
(2,'f'): glUniform2fv, | |
(3,'f'): glUniform3fv, | |
(4,'f'): glUniform4fv, | |
(2,2): glUniformMatrix2fv, | |
(2,3): glUniformMatrix2x3fv, | |
(2,4): glUniformMatrix2x4fv, | |
(3,2): glUniformMatrix3x2fv, | |
(3,3): glUniformMatrix3fv, | |
(3,4): glUniformMatrix3x4fv, | |
(4,2): glUniformMatrix4x2fv, | |
(4,3): glUniformMatrix4x3fv, | |
(4,4): glUniformMatrix4fv, | |
} | |