popsicle/examples/webgpu/pbr2.py at master · kunitoki/popsicle · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
"""

PBR Rendering 2
===============

This example shows the lighting rendering affect of materials with different
metalness and roughness. Every second sphere has an IBL environment map on it.
"""

# run_example = false
# sphinx_gallery_pygfx_animate = True
# sphinx_gallery_pygfx_duration = 3
# sphinx_gallery_pygfx_lossless = False

import math
from colorsys import hls_to_rgb
from itertools import count

import imageio.v3 as iio
import numpy as np
from pop import WgpuCanvas, run

import pygfx as gfx
import pylinalg as la


def main(canvas):
    frame_idx = count()  # counter for animation
    renderer = gfx.renderers.WgpuRenderer(canvas)
    scene = gfx.Scene()

    # Lights
    scene.add(gfx.AmbientLight("#fff", 0.2))
    directional_light = gfx.DirectionalLight("#fff", 3)
    directional_light.local.position = la.vec_normalize((1, 1, 1))
    scene.add(directional_light)
    point_light = gfx.PointLight("#fff", 3)
    scene.add(point_light)
    point_light.add(gfx.PointLightHelper(4))

    # Read cube image and turn it into a 3D image (a 4d array)
    env_img = iio.imread("imageio:meadow_cube.jpg")
    cube_size = env_img.shape[1]
    env_img.shape = 6, cube_size, cube_size, env_img.shape[-1]

    # Create environment map
    env_tex = gfx.Texture(
        env_img, dim=2, size=(cube_size, cube_size, 6), generate_mipmaps=True
    )

    # Apply env map to skybox
    background = gfx.Background(None, gfx.BackgroundSkyboxMaterial(map=env_tex))
    scene.add(background)

    # Now create spheres ...
    cube_width = 400
    numbers_per_side = 5
    sphere_radius = (cube_width / numbers_per_side) * 0.8 * 0.5
    step_size = 1.0 / numbers_per_side

    geometry = gfx.sphere_geometry(sphere_radius, 32, 16)

    index = 0
    alpha = 0.0
    while alpha <= 1.0:
        beta = 0.0
        while beta <= 1.0:
            gamma = 0.0
            while gamma <= 1.0:
                material = gfx.MeshStandardMaterial(
                    color=hls_to_rgb(alpha, 0.5, gamma * 0.5 + 0.1),
                    metalness=beta,
                    roughness=1.0 - alpha,
                )

                if index % 2 != 0:
                    material.env_map = env_tex

                mesh = gfx.Mesh(geometry, material)

                mesh.local.position = (
                    alpha * 400 - 200,
                    beta * 400 - 200,
                    gamma * 400 - 200,
                )
                scene.add(mesh)
                index += 1

                gamma += step_size
            beta += step_size
        alpha += step_size

    # Create camera and controls
    camera = gfx.PerspectiveCamera(45, 640 / 480)
    camera.show_object(scene, view_dir=(-2, -1.5, -3), scale=1.2)

    controller = gfx.OrbitController(camera, register_events=renderer)

    def animate():
        timer = next(frame_idx) / 30

        point_light.local.position = (
            math.sin(timer / 30 * (2 * np.pi)) * 300,
            math.cos(timer * 2 / 30 * (2 * np.pi)) * 400,
            math.cos(timer / 30 * (2 * np.pi)) * 300,
        )

        renderer.render(scene, camera)
        renderer.request_draw()

    renderer.request_draw(animate)


if __name__ == "__main__":
    canvas = WgpuCanvas(size=(640, 480), title="gfx_pbr")
    run(canvas, main)