motionlab/shaders.py

"""Shader layers: every effects/*.glsl file becomes a layer, reloaded whenever the file is saved.

A shader file only contains its own sliders and a main(). The host supplies the rest:

    in  vec2 uv;                 pixel position, (0,0) top-left .. (1,1) bottom-right
    out vec4 color;              what to write for this pixel
    uniform sampler2D camera;    the incoming frame (output of the layer above), RGB
    uniform sampler2D prev;      this layer's own output from the previous frame
    uniform sampler2D flow;      .xy = how far this pixel moved since last frame, in uv units
    uniform sampler2D motion;    .r  = how much this pixel changed since last frame, 0..1
    uniform float time;          seconds since start
    uniform float energy;        overall amount of motion in the frame, 0..1
    uniform vec2 resolution;     output size in pixels

Sliders are plain float uniforms with a range comment:

    uniform float amount;  // 0.0 .. 0.05 = 0.01
"""
import os
import re

import moderngl
import numpy as np

from effects import Layer, Param

VERTEX = """#version 410
out vec2 uv;
void main() {
    uv = vec2((gl_VertexID << 1) & 2, gl_VertexID & 2);
    gl_Position = vec4(uv * 2.0 - 1.0, 0.0, 1.0);
}
"""

# "#line 1" makes compiler errors report line numbers of the user's file
HEADER = """#version 410
uniform sampler2D camera;
uniform sampler2D prev;
uniform sampler2D flow;
uniform sampler2D motion;
uniform float time;
uniform float energy;
uniform vec2 resolution;
in vec2 uv;
out vec4 color;
#line 1
"""

NUM = r"(-?[\d.]+)"
PARAM_RE = re.compile(rf"uniform\s+float\s+(\w+)\s*;(?:[ \t]*//\s*{NUM}\s*\.\.\s*{NUM}\s*=\s*{NUM})?")


class ShaderLayer(Layer):
    enabled = False

    def __init__(self, gl, path):
        super().__init__()
        self.gl, self.path = gl, path
        self.name = os.path.splitext(os.path.basename(path))[0]
        self.program = None
        self.error = None
        self.mtime = None
        self.targets = None
        self.reload()

    def reload(self):
        """Recompile if the file changed. On a compile error the last working version keeps running."""
        try:
            mtime = os.stat(self.path).st_mtime
        except FileNotFoundError:
            return
        if mtime == self.mtime:
            return
        self.mtime = mtime
        with open(self.path) as f:
            source = f.read()
        try:
            program = self.gl.program(vertex_shader=VERTEX, fragment_shader=HEADER + source)
        except moderngl.Error as e:
            self.error = str(e)
            print(f"[{self.name}] shader error:\n{e}", flush=True)
            return
        self.error = None
        self.program = program
        self.vao = self.gl.vertex_array(program, [])
        old = self.params
        self.params = {name: Param(float(default or 0.5), float(lo or 0), float(hi or 1))
                       for name, lo, hi, default in PARAM_RE.findall(source)}
        for name, p in self.params.items():  # keep slider positions across reloads
            if name in old:
                p.value = min(max(old[name].value, p.lo), p.hi)
        print(f"[{self.name}] loaded", flush=True)

    def _target(self, size):
        texture = self.gl.texture(size, 4)
        texture.repeat_x = texture.repeat_y = False
        return texture, self.gl.framebuffer([texture])

    def _texture(self, attr, size, components, dtype="f1"):
        if not hasattr(self, attr):
            texture = self.gl.texture(size, components, dtype=dtype)
            texture.repeat_x = texture.repeat_y = False
            setattr(self, attr, texture)
        return getattr(self, attr)

    def _set(self, name, value):
        uniform = self.program.get(name, None)  # absent if the shader doesn't use it
        if uniform is not None:
            uniform.value = value
        return uniform is not None

    def apply(self, frame, ctx):
        if self.program is None:
            return frame
        h, w = frame.shape[:2]
        if self.targets is None:
            # two render targets, swapped every frame, so the shader can read its previous output
            self.targets = [self._target((w, h)), self._target((w, h))]
        (_, framebuffer), (prev, _) = self.targets

        camera = self._texture("camera", (w, h), 3)
        camera.swizzle = "BGR1"  # frames are BGR; shaders see RGB
        camera.write(np.ascontiguousarray(frame))
        camera.use(0)
        self._set("camera", 0)
        prev.use(1)
        self._set("prev", 1)
        if self._set("flow", 2):
            fh, fw = ctx.flow.shape[:2]
            flow = self._texture("flow_tex", (fw, fh), 2, "f4")
            flow.write((ctx.flow / (fw, fh)).astype("f4"))
            flow.use(2)
        if self._set("motion", 3):
            mh, mw = ctx.motion.shape
            motion = self._texture("motion_tex", (mw, mh), 1)
            motion.write(ctx.motion)
            motion.use(3)
        self._set("time", ctx.t)
        self._set("energy", ctx.energy)
        self._set("resolution", (w, h))
        for name, p in self.params.items():
            self._set(name, p.value)

        framebuffer.use()
        self.vao.render(moderngl.TRIANGLES, vertices=3)
        out = np.frombuffer(framebuffer.read(components=3), np.uint8).reshape(h, w, 3)
        self.targets.reverse()
        return np.ascontiguousarray(out[:, :, ::-1])


class ShaderBank:
    """Watches a folder of .glsl files and keeps one ShaderLayer per file in the shared layer list."""

    def __init__(self, directory, layers):
        self.directory, self.layers = directory, layers
        self.known = {}
        self.gl = moderngl.create_context(standalone=True, require=410)
        self.refresh()

    def refresh(self):
        if os.path.isdir(self.directory):
            for filename in sorted(os.listdir(self.directory)):
                path = os.path.join(self.directory, filename)
                if filename.endswith(".glsl") and path not in self.known:
                    self.known[path] = ShaderLayer(self.gl, path)
                    self.layers.append(self.known[path])
        for layer in self.known.values():
            layer.reload()

    @property
    def errors(self):
        return [layer.name for layer in self.known.values() if layer.error]