Skip to content

Commit

Permalink
Improve Shadertoy Utility (#401)
Browse files Browse the repository at this point in the history
* 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 <[email protected]>
  • Loading branch information
Vipitis and Korijn authored Nov 6, 2023
1 parent 6ce9ac9 commit dbb991d
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 4 deletions.
87 changes: 87 additions & 0 deletions tests/test_util_shadertoy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
43 changes: 39 additions & 4 deletions wgpu/utils/shadertoy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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),
Expand All @@ -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()

Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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__":
Expand Down

0 comments on commit dbb991d

Please sign in to comment.