From dbb991d23f31bbbb568bbe3e6d8cb797e0f111a2 Mon Sep 17 00:00:00 2001 From: Jan Date: Mon, 6 Nov 2023 17:59:05 +0100 Subject: [PATCH] Improve Shadertoy Utility (#401) * Add offscreen support * Add snapshot method * Add shadercode validation for glsl * Remove unneeded dependencies * Add additional validation with wgsl translation * Add tests for offscreen and snapshot * Revert external validation * Fix style * Fix docstring typos --------- Co-authored-by: Korijn van Golen --- tests/test_util_shadertoy.py | 87 ++++++++++++++++++++++++++++++++++++ wgpu/utils/shadertoy.py | 43 ++++++++++++++++-- 2 files changed, 126 insertions(+), 4 deletions(-) diff --git a/tests/test_util_shadertoy.py b/tests/test_util_shadertoy.py index d1372f93..7429d484 100644 --- a/tests/test_util_shadertoy.py +++ b/tests/test_util_shadertoy.py @@ -66,3 +66,90 @@ def test_shadertoy_glsl(): assert shader.shader_type == "glsl" shader._draw_frame() + + +def test_shadertoy_offscreen(): + # Import here, because it imports the wgpu.gui.auto + from wgpu.utils.shadertoy import Shadertoy # noqa + + shader_code = """ + void shader_main(out vec4 fragColor, vec2 frag_coord) { + vec2 uv = frag_coord / i_resolution.xy; + + if ( length(frag_coord - i_mouse.xy) < 20.0 ) { + fragColor = vec4(0.0, 0.0, 0.0, 1.0); + }else{ + fragColor = vec4( 0.5 + 0.5 * sin(i_time * vec3(uv, 1.0) ), 1.0); + } + + } + """ + + shader = Shadertoy(shader_code, resolution=(800, 450), offscreen=True) + assert shader.resolution == (800, 450) + assert shader.shader_code == shader_code + assert shader.shader_type == "glsl" + assert shader._offscreen is True + + +def test_shadertoy_snapshot(): + # Import here, because it imports the wgpu.gui.auto + from wgpu.utils.shadertoy import Shadertoy # noqa + + shader_code = """ + void shader_main(out vec4 fragColor, vec2 frag_coord) { + vec2 uv = frag_coord / i_resolution.xy; + + if ( length(frag_coord - i_mouse.xy) < 20.0 ) { + fragColor = vec4(0.0, 0.0, 0.0, 1.0); + }else{ + fragColor = vec4( 0.5 + 0.5 * sin(i_time * vec3(uv, 1.0) ), 1.0); + } + + } + """ + + shader = Shadertoy(shader_code, resolution=(800, 450), offscreen=True) + frame1a = shader.snapshot( + time_float=0.0, + mouse_pos=( + 0, + 0, + 0, + 0, + ), + ) + frame2a = shader.snapshot( + time_float=1.2, + mouse_pos=( + 100, + 200, + 0, + 0, + ), + ) + frame1b = shader.snapshot( + time_float=0.0, + mouse_pos=( + 0, + 0, + 0, + 0, + ), + ) + frame2b = shader.snapshot( + time_float=1.2, + mouse_pos=( + 100, + 200, + 0, + 0, + ), + ) + + assert shader.resolution == (800, 450) + assert shader.shader_code == shader_code + assert shader.shader_type == "glsl" + assert shader._offscreen is True + assert frame1a == frame1b + assert frame2a == frame2b diff --git a/wgpu/utils/shadertoy.py b/wgpu/utils/shadertoy.py index 849c55b3..3ab36fd9 100644 --- a/wgpu/utils/shadertoy.py +++ b/wgpu/utils/shadertoy.py @@ -3,7 +3,7 @@ import wgpu from wgpu.gui.auto import WgpuCanvas, run - +from wgpu.gui.offscreen import WgpuCanvas as OffscreenCanvas, run as run_offscreen vertex_code_glsl = """ #version 450 core @@ -222,6 +222,7 @@ class Shadertoy: Parameters: shader_code (str): The shader code to use. resolution (tuple): The resolution of the shadertoy. + offscreen (bool): Whether to render offscreen. Default is False. The shader code must contain a entry point function: @@ -247,7 +248,7 @@ class Shadertoy: # todo: support input textures # todo: support multiple render passes (`i_channel0`, `i_channel1`, etc.) - def __init__(self, shader_code, resolution=(800, 450)) -> None: + def __init__(self, shader_code, resolution=(800, 450), offscreen=False) -> None: self._uniform_data = UniformArray( ("mouse", "f", 4), ("resolution", "f", 3), @@ -259,6 +260,8 @@ def __init__(self, shader_code, resolution=(800, 450)) -> None: self._shader_code = shader_code self._uniform_data["resolution"] = resolution + (1,) + self._offscreen = offscreen + self._prepare_render() self._bind_events() @@ -288,7 +291,14 @@ def shader_type(self): def _prepare_render(self): import wgpu.backends.rs # noqa - self._canvas = WgpuCanvas(title="Shadertoy", size=self.resolution, max_fps=60) + if self._offscreen: + self._canvas = OffscreenCanvas( + title="Shadertoy", size=self.resolution, max_fps=60 + ) + else: + self._canvas = WgpuCanvas( + title="Shadertoy", size=self.resolution, max_fps=60 + ) adapter = wgpu.request_adapter( canvas=self._canvas, power_preference="high-performance" @@ -463,7 +473,32 @@ def _draw_frame(self): def show(self): self._canvas.request_draw(self._draw_frame) - run() + if self._offscreen: + run_offscreen() + else: + run() + + def snapshot(self, time_float: float = 0.0, mouse_pos: tuple = (0, 0, 0, 0)): + """ + Returns an image of the specified time. (Only available when ``offscreen=True``) + + Parameters: + time_float (float): The time to snapshot. It essentially sets ``i_time`` to a specific number. (Default is 0.0) + mouse_pos (tuple): The mouse position in pixels in the snapshot. It essentially sets ``i_mouse`` to a 4-tuple. (Default is (0,0,0,0)) + Returns: + frame (memoryview): snapshot with transparancy. This object can be converted to a numpy array (without copying data) + using ``np.asarray(arr)`` + """ + if not self._offscreen: + raise NotImplementedError("Snapshot is only available in offscreen mode.") + + if hasattr(self, "_last_time"): + self.__delattr__("_last_time") + self._uniform_data["time"] = time_float + self._uniform_data["mouse"] = mouse_pos + self._canvas.request_draw(self._draw_frame) + frame = self._canvas.draw() + return frame if __name__ == "__main__":