LutGrading

Desktop 3D-LUT colour grading via WorldEnvironment.

▶ Run in browser

Tags: 3d post-process colour-grading lut tonemap

Drives WorldEnvironment.lut_enabled / lut_tex_id on the desktop Vulkan renderer, matching the web backend: a 3D rgba8 LUT sampled post-tonemap on the final LDR colour. A warm grade (id 1) and a cool/teal grade (id 2) are registered at startup; pressing a key swaps the active LUT or disables grading entirely. The identity (no-op) state is lut_enabled = False.

Controls: A / D - Orbit camera left / right W / S - Pitch camera up / down 1 - No LUT (neutral) 2 - Warm grade 3 - Cool / teal grade Escape - Quit

Source

  1"""LutGrading: Desktop 3D-LUT colour grading via WorldEnvironment.
  2
  3# /// simvx
  4# tags = ["post-process", "colour-grading", "lut", "tonemap"]
  5# web = { width = 1280, height = 720 }
  6# ///
  7
  8Drives ``WorldEnvironment.lut_enabled`` / ``lut_tex_id`` on the desktop Vulkan
  9renderer, matching the web backend: a 3D ``rgba8`` LUT sampled post-tonemap on
 10the final LDR colour. A warm grade (id 1) and a cool/teal grade (id 2) are
 11registered at startup; pressing a key swaps the active LUT or disables grading
 12entirely. The identity (no-op) state is ``lut_enabled = False``.
 13
 14Controls:
 15    A / D       - Orbit camera left / right
 16    W / S       - Pitch camera up / down
 17    1           - No LUT (neutral)
 18    2           - Warm grade
 19    3           - Cool / teal grade
 20    Escape      - Quit
 21"""
 22
 23import math
 24import sys
 25
 26import numpy as np
 27
 28from simvx.core import (
 29    Camera3D,
 30    DirectionalLight3D,
 31    Input,
 32    InputMap,
 33    Key,
 34    Material,
 35    Mesh,
 36    MeshInstance3D,
 37    Node,
 38    Text2D,
 39    WorldEnvironment,
 40)
 41from simvx.core.colour_grading import generate_cool_lut, generate_warm_lut
 42from simvx.graphics import App
 43
 44LUT_WARM = 1
 45LUT_COOL = 2
 46
 47
 48class LutGrading(Node):
 49    def on_ready(self):
 50        InputMap.add_action("orbit_left", [Key.A])
 51        InputMap.add_action("orbit_right", [Key.D])
 52        InputMap.add_action("pitch_up", [Key.W])
 53        InputMap.add_action("pitch_down", [Key.S])
 54        InputMap.add_action("lut_off", [Key.KEY_1])
 55        InputMap.add_action("lut_warm", [Key.KEY_2])
 56        InputMap.add_action("lut_cool", [Key.KEY_3])
 57        InputMap.add_action("quit", [Key.ESCAPE])
 58
 59        self._yaw = 35.0
 60        self._pitch = 22.0
 61        self._distance = 16.0
 62        self._target = (0.0, 1.0, 0.0)
 63
 64        self._cam = Camera3D(name="Camera", fov=55, near=0.1, far=200.0)
 65        self.add_child(self._cam)
 66
 67        env = self.add_child(WorldEnvironment())
 68        env.tonemap_mode = "aces"
 69        env.tonemap_white = 1.0
 70        env.tonemap_exposure = 1.0
 71        env.sky_mode = "colour"
 72        env.lut_enabled = False
 73        env.lut_tex_id = 0
 74        self._env = env
 75
 76        key = DirectionalLight3D(name="KeyLight", intensity=2.2)
 77        key.look_at((-0.6, -1.0, -0.4))
 78        self.add_child(key)
 79        fill = DirectionalLight3D(name="FillLight", intensity=0.3, colour=(0.7, 0.8, 1.0))
 80        fill.look_at((0.8, -0.6, 1.0))
 81        self.add_child(fill)
 82
 83        ground = MeshInstance3D(name="Ground", mesh=Mesh.cube())
 84        ground.material = Material(colour=(0.6, 0.6, 0.62), roughness=0.85, metallic=0.0)
 85        ground.scale = (40.0, 0.1, 40.0)
 86        ground.position = (0.0, -0.05, 0.0)
 87        self.add_child(ground)
 88
 89        palette = [
 90            (0.85, 0.85, 0.85),
 91            (0.8, 0.8, 0.82),
 92            (0.75, 0.78, 0.85),
 93            (0.82, 0.8, 0.78),
 94            (0.78, 0.82, 0.8),
 95        ]
 96        rng = np.random.default_rng(3)
 97        for i in range(20):
 98            colour = palette[i % len(palette)]
 99            mat = Material(colour=colour, roughness=0.6, metallic=0.05)
100            mesh = Mesh.cube() if i % 2 == 0 else Mesh.sphere(radius=0.6)
101            obj = MeshInstance3D(name=f"Obj{i}", mesh=mesh, material=mat)
102            ring = 3.0 + (i % 3) * 2.5
103            angle = i * math.pi * 2 / 7
104            obj.position = (
105                math.cos(angle) * ring + rng.uniform(-0.4, 0.4),
106                0.6 + (i % 3) * 0.3,
107                math.sin(angle) * ring + rng.uniform(-0.4, 0.4),
108            )
109            self.add_child(obj)
110
111        # Register the grading LUTs once the renderer (and its post-process pass)
112        # is live. ``app`` is available from on_ready onward.
113        self.app.register_lut(LUT_WARM, generate_warm_lut(32))
114        self.app.register_lut(LUT_COOL, generate_cool_lut(32))
115
116        self._hud = self.add_child(Text2D(name="HUD", text="", font_scale=1.8, x=12.0, y=12.0))
117        self._update_camera()
118        self._update_hud()
119
120    def on_process(self, dt):
121        if Input.is_action_just_pressed("quit"):
122            self.app.quit()
123            return
124
125        if Input.is_action_pressed("orbit_left"):
126            self._yaw += 60.0 * dt
127        if Input.is_action_pressed("orbit_right"):
128            self._yaw -= 60.0 * dt
129        if Input.is_action_pressed("pitch_up"):
130            self._pitch = min(80.0, self._pitch + 30.0 * dt)
131        if Input.is_action_pressed("pitch_down"):
132            self._pitch = max(-10.0, self._pitch - 30.0 * dt)
133
134        env = self._env
135        if Input.is_action_just_pressed("lut_off"):
136            env.lut_enabled = False
137            env.lut_tex_id = 0
138        if Input.is_action_just_pressed("lut_warm"):
139            env.lut_tex_id = LUT_WARM
140            env.lut_enabled = True
141        if Input.is_action_just_pressed("lut_cool"):
142            env.lut_tex_id = LUT_COOL
143            env.lut_enabled = True
144
145        self._update_camera()
146        self._update_hud()
147
148    def _update_camera(self):
149        yaw_rad = math.radians(self._yaw)
150        pitch_rad = math.radians(self._pitch)
151        cp = math.cos(pitch_rad)
152        x = self._target[0] + self._distance * cp * math.sin(yaw_rad)
153        y = self._target[1] + self._distance * math.sin(pitch_rad)
154        z = self._target[2] + self._distance * cp * math.cos(yaw_rad)
155        self._cam.position = (x, y, z)
156        self._cam.look_at(self._target)
157
158    def _update_hud(self):
159        env = self._env
160        if not env.lut_enabled:
161            grade = "none (neutral)"
162        elif env.lut_tex_id == LUT_WARM:
163            grade = "warm (amber)"
164        elif env.lut_tex_id == LUT_COOL:
165            grade = "cool (teal)"
166        else:
167            grade = f"id {env.lut_tex_id}"
168        self._hud.text = "\n".join(
169            [
170                "3D LUT Colour Grading (WorldEnvironment -> desktop)",
171                f"Active grade: {grade}",
172                "[1] none  [2] warm  [3] cool   A/D orbit  W/S pitch  Esc quit",
173            ]
174        )
175
176
177def _avg_rgb(frame: np.ndarray) -> np.ndarray:
178    """Mean linear RGB of a captured RGBA frame, as float in [0, 1]."""
179    return frame[..., :3].astype(np.float32).mean(axis=(0, 1)) / 255.0
180
181
182def main() -> None:
183    if "--test" in sys.argv:
184        from simvx.graphics.testing import assert_not_blank
185
186        # Render the same scene twice: once neutral, once with the warm LUT, and
187        # assert the warm grade shifts the average colour toward red / away from
188        # blue (the warm LUT boosts R, cuts B). Both must be non-blank.
189        app = App(title="LutGrading", width=640, height=480, visible=False)
190
191        class _Neutral(LutGrading):
192            def on_ready(self):
193                super().on_ready()
194                self._env.lut_enabled = False
195
196        neutral_frames = app.run_headless(_Neutral(), frames=4, capture_frames=[3])
197        assert neutral_frames, "no neutral frame captured"
198        assert_not_blank(neutral_frames[0])
199        neutral_avg = _avg_rgb(neutral_frames[0])
200
201        class _Warm(LutGrading):
202            def on_ready(self):
203                super().on_ready()
204                self._env.lut_tex_id = LUT_WARM
205                self._env.lut_enabled = True
206
207        warm_frames = app.run_headless(_Warm(), frames=4, capture_frames=[3])
208        assert warm_frames, "no warm frame captured"
209        assert_not_blank(warm_frames[0])
210        warm_avg = _avg_rgb(warm_frames[0])
211
212        # Warm grade: red up relative to blue versus neutral.
213        neutral_rb = float(neutral_avg[0] - neutral_avg[2])
214        warm_rb = float(warm_avg[0] - warm_avg[2])
215        assert (
216            warm_rb > neutral_rb + 0.02
217        ), f"warm LUT did not shift colour warm: neutral R-B={neutral_rb:.3f}, warm R-B={warm_rb:.3f}"
218        print(
219            f"lut_grading --test OK  neutral_avg={neutral_avg.round(3).tolist()} "
220            f"warm_avg={warm_avg.round(3).tolist()}  R-B neutral={neutral_rb:.3f} warm={warm_rb:.3f}"
221        )
222        return
223
224    app = App(title="LutGrading", width=1280, height=720)
225    app.run(LutGrading())
226
227
228if __name__ == "__main__":
229    main()