GlintsPlayground/glints.py at main · WeakKnight/GlintsPlayground · 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
# SPDX-License-Identifier: Apache-2.0

import slangpy as spy
import sgl
from pathlib import Path
import numpy as np
from slangpy.types import call_id
import math
import imageio
import os

device = spy.create_device(include_paths = [Path(__file__).parent])
module = spy.Module.load_from_file(device, "glints.slang")

noiseModule = spy.Module.load_from_file(device, "glints2023Noise.slang")
noiseSize = 512
noiseTex = device.create_texture(
    width = noiseSize,
    height = noiseSize,
    format = sgl.Format.rgba32_float,
    usage = sgl.ResourceUsage.shader_resource | sgl.ResourceUsage.unordered_access
)
noiseModule.InitNoise(call_id(), sgl.float2(noiseSize, noiseSize), 42, _result = noiseTex)

output_w = 1920
output_h = 1080

uniforms = {
    "_type": "Uniforms",
    "screenSize": sgl.float2(output_w, output_h),
    "focalLength": 24.0,
    "frameHeight": 24.0,
}

radius = 15.0
alpha: float = 0.0
pi: float = 3.14159265359
beta: float = -pi / 4.0

def updateCamera():
    global alpha, beta, pi
    cameraDir = [
        -math.cos(alpha) * math.sin(beta),
        -math.cos(beta),
        -math.sin(alpha) * math.sin(beta),
    ]
    betaUp = beta + pi * 0.5
    cameraUp = [
        math.cos(alpha) * math.sin(betaUp),
        math.cos(betaUp),
        math.sin(alpha) * math.sin(betaUp),
    ]
    cameraRight = np.cross(cameraDir, cameraUp).tolist()
    cameraPos = [-cameraDir[0] * radius, -cameraDir[1] * radius, -cameraDir[2] * radius]
    uniforms["cameraDir"] = cameraDir
    uniforms["cameraUp"] = cameraUp
    uniforms["cameraRight"] = cameraRight
    uniforms["cameraPosition"] = cameraPos
    return

updateCamera()

pathname = os.path.realpath(__file__)

def loadImageData(path, w, h):
    imageData = imageio.v3.imread(path)
    # reshape imageData into 512x512x4, adding an alpha channel
    print(imageData.shape)
    if len(imageData.shape) >= 3 and imageData.shape[2] == 3:
        imageData = np.concatenate([imageData, np.ones((w, h, 1), dtype=np.uint8) * 255], axis=2)
    return imageData

baseColorImageData = loadImageData("albedo.jpg", 4096, 4096)
normalImageData = loadImageData("normal.jpg", 4096, 4096)
roughnessImageData = loadImageData("roughness.jpg", 4096, 4096)

albedoTex = device.create_texture(
    width = 4096,
    height = 4096,
    format = sgl.Format.rgba8_unorm_srgb,
    usage = sgl.ResourceUsage.shader_resource,
    data = baseColorImageData,
)

normalTex = device.create_texture(
    width = 4096,
    height = 4096,
    format = sgl.Format.rgba8_unorm,
    usage = sgl.ResourceUsage.shader_resource,
    data = normalImageData,
)

roughnessTex = device.create_texture(
    width = 4096,
    height = 4096,
    format = sgl.Format.r8_unorm,
    usage = sgl.ResourceUsage.shader_resource,
    data = roughnessImageData,
)

samplerState = device.create_sampler()

output = device.create_texture(
    width = output_w,
    height = output_h,
    format = sgl.Format.rgba8_unorm,
    usage = sgl.ResourceUsage.shader_resource | sgl.ResourceUsage.unordered_access
)
module.raytraceScene(call_id(), uniforms, albedoTex, normalTex, roughnessTex, noiseTex, samplerState, _result = output)

imageio.imwrite("out.png", output.to_numpy())