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, | |
| } | |