hexdoc.graphics.render

  1# pyright: reportUnknownMemberType=false
  2
  3from __future__ import annotations
  4
  5import logging
  6import math
  7from dataclasses import dataclass
  8from enum import Flag, auto
  9from functools import cached_property
 10from pathlib import Path
 11from typing import Any, Literal, cast
 12
 13import importlib_resources as resources
 14import moderngl as mgl
 15import moderngl_window as mglw
 16import numpy as np
 17from moderngl import Context, Program, Uniform
 18from moderngl_window import WindowConfig
 19from moderngl_window.context.headless import Window as HeadlessWindow
 20from moderngl_window.opengl.vao import VAO
 21from PIL import Image
 22from pydantic import ValidationError
 23from pyrr import Matrix44
 24
 25from hexdoc.core import ModResourceLoader, ResourceLocation
 26from hexdoc.graphics import glsl
 27from hexdoc.minecraft.assets import AnimationMeta
 28from hexdoc.minecraft.assets.animated import AnimationMetaTag
 29from hexdoc.minecraft.models import BlockModel
 30from hexdoc.minecraft.models.base_model import (
 31    DisplayPosition,
 32    ElementFace,
 33    ElementFaceUV,
 34    FaceName,
 35    ModelElement,
 36)
 37from hexdoc.utils.types import Vec3, Vec4
 38
 39logger = logging.getLogger(__name__)
 40
 41
 42# https://minecraft.wiki/w/Help:Isometric_renders#Preferences
 43LIGHT_TOP = 0.98
 44LIGHT_LEFT = 0.8
 45LIGHT_RIGHT = 0.608
 46
 47LIGHT_FLAT = 0.98
 48
 49
 50class DebugType(Flag):
 51    NONE = 0
 52    AXES = auto()
 53    NORMALS = auto()
 54
 55
 56@dataclass(kw_only=True)
 57class BlockRenderer:
 58    loader: ModResourceLoader
 59    output_dir: Path | None = None
 60    debug: DebugType = DebugType.NONE
 61
 62    def __post_init__(self):
 63        self.window = HeadlessWindow(
 64            size=(300, 300),
 65        )
 66        mglw.activate_context(self.window)
 67
 68        self.config = BlockRendererConfig(ctx=self.window.ctx, wnd=self.window)
 69
 70        self.window.config = self.config
 71        self.window.swap_buffers()
 72        self.window.set_default_viewport()
 73
 74    @property
 75    def ctx(self):
 76        return self.window.ctx
 77
 78    def render_block_model(
 79        self,
 80        model: BlockModel | ResourceLocation,
 81        output_path: str | Path,
 82    ):
 83        if isinstance(model, ResourceLocation):
 84            _, model = self.loader.load_resource(
 85                type="assets",
 86                folder="models",
 87                id=model,
 88                decode=BlockModel.model_validate_json,
 89            )
 90
 91        model.load_parents_and_apply(self.loader)
 92
 93        textures = {
 94            name: self.load_texture(texture_id)
 95            for name, texture_id in model.resolve_texture_variables().items()
 96        }
 97
 98        output_path = Path(output_path)
 99        if self.output_dir and not output_path.is_absolute():
100            output_path = self.output_dir / output_path
101
102        self.config.render_block(model, textures, output_path, self.debug)
103
104    def load_texture(self, texture_id: ResourceLocation):
105        logger.debug(f"Loading texture: {texture_id}")
106        _, path = self.loader.find_resource("assets", "textures", texture_id + ".png")
107
108        meta_path = path.with_suffix(".png.mcmeta")
109        if meta_path.is_file():
110            logger.debug(f"Loading animation mcmeta: {meta_path}")
111            # FIXME: hack
112            try:
113                meta = AnimationMeta.model_validate_json(meta_path.read_bytes())
114            except ValidationError as e:
115                logger.warning(f"Failed to parse animation meta for {texture_id}:\n{e}")
116                meta = None
117        else:
118            meta = None
119
120        return BlockTextureInfo(path, meta)
121
122    def destroy(self):
123        self.window.destroy()
124
125    def __enter__(self):
126        return self
127
128    def __exit__(self, *_: Any):
129        self.destroy()
130        return False
131
132
133class BlockRendererConfig(WindowConfig):
134    def __init__(self, ctx: Context, wnd: HeadlessWindow):
135        super().__init__(ctx, wnd)
136
137        # depth test: ensure faces are displayed in the correct order
138        # blend: handle translucency
139        # cull face: remove back faces, eg. for trapdoors
140        self.ctx.enable(mgl.DEPTH_TEST | mgl.BLEND | mgl.CULL_FACE)
141
142        view_size = 16
143        self.projection = Matrix44.orthogonal_projection(
144            left=-view_size / 2,
145            right=view_size / 2,
146            top=view_size / 2,
147            bottom=-view_size / 2,
148            near=0.01,
149            far=20_000,
150            dtype="f4",
151        ) * Matrix44.from_scale((1, -1, 1), "f4")
152
153        self.camera, self.eye = direction_camera(pos="south")
154
155        self.lights = [
156            ((0, -1, 0), LIGHT_TOP),
157            ((1, 0, -1), LIGHT_LEFT),
158            ((-1, 0, -1), LIGHT_RIGHT),
159        ]
160
161        # block faces
162
163        self.face_prog = self.ctx.program(
164            vertex_shader=read_shader("block_face", "vertex"),
165            fragment_shader=read_shader("block_face", "fragment"),
166        )
167
168        self.uniform("m_proj").write(self.projection)
169        self.uniform("m_camera").write(self.camera)
170        self.uniform("layer").value = 0  # TODO: implement animations
171
172        for i, (direction, diffuse) in enumerate(self.lights):
173            self.uniform(f"lights[{i}].direction").value = direction
174            self.uniform(f"lights[{i}].diffuse").value = diffuse
175
176        # axis planes
177
178        self.debug_plane_prog = self.ctx.program(
179            vertex_shader=read_shader("debug/plane", "vertex"),
180            fragment_shader=read_shader("debug/plane", "fragment"),
181        )
182
183        self.uniform("m_proj", self.debug_plane_prog).write(self.projection)
184        self.uniform("m_camera", self.debug_plane_prog).write(self.camera)
185        self.uniform("m_model", self.debug_plane_prog).write(Matrix44.identity("f4"))
186
187        self.debug_axes = list[tuple[VAO, Vec4]]()
188
189        pos = 8
190        neg = 0
191        for from_, to, color, direction in [
192            ((0, neg, neg), (0, pos, pos), (1, 0, 0, 0.75), "east"),
193            ((neg, 0, neg), (pos, 0, pos), (0, 1, 0, 0.75), "up"),
194            ((neg, neg, 0), (pos, pos, 0), (0, 0, 1, 0.75), "south"),
195        ]:
196            vao = VAO()
197            verts = get_face_verts(from_, to, direction)
198            vao.buffer(np.array(verts, np.float32), "3f", ["in_position"])
199            self.debug_axes.append((vao, color))
200
201        # vertex normal vectors
202
203        self.debug_normal_prog = self.ctx.program(
204            vertex_shader=read_shader("debug/normal", "vertex"),
205            geometry_shader=read_shader("debug/normal", "geometry"),
206            fragment_shader=read_shader("debug/normal", "fragment"),
207        )
208
209        self.uniform("m_proj", self.debug_normal_prog).write(self.projection)
210        self.uniform("m_camera", self.debug_normal_prog).write(self.camera)
211        self.uniform("lineSize", self.debug_normal_prog).value = 4
212
213        self.ctx.line_width = 3
214
215    def render_block(
216        self,
217        model: BlockModel,
218        texture_vars: dict[str, BlockTextureInfo],
219        output_path: Path,
220        debug: DebugType = DebugType.NONE,
221    ):
222        if not model.elements:
223            raise ValueError("Unable to render model, no elements found")
224
225        self.wnd.clear()
226
227        # enable/disable flat item lighting
228        match model.gui_light:
229            case "front":
230                flatLighting = LIGHT_FLAT
231            case "side":
232                flatLighting = 0
233        self.uniform("flatLighting").value = flatLighting
234
235        # load textures
236        texture_locs = dict[str, int]()
237        transparent_textures = set[str]()
238
239        for i, (name, info) in enumerate(texture_vars.items()):
240            texture_locs[name] = i
241
242            logger.debug(f"Loading texture {name}: {info}")
243            image = Image.open(info.image_path).convert("RGBA")
244
245            extrema = image.getextrema()
246            assert len(extrema) >= 4, f"Expected 4 bands but got {len(extrema)}"
247            min_alpha, _ = extrema[3]
248            if min_alpha < 255:
249                logger.debug(f"Transparent texture: {name} ({min_alpha=})")
250                transparent_textures.add(name)
251
252            # TODO: implement non-square animations, write test cases
253            match info.meta:
254                case AnimationMeta(
255                    animation=AnimationMetaTag(height=frame_height),
256                ) if frame_height:
257                    # animated with specified size
258                    layers = image.height // frame_height
259                case AnimationMeta():
260                    # size is unspecified, assume it's square and verify later
261                    frame_height = image.width
262                    layers = image.height // frame_height
263                case None:
264                    # non-animated
265                    frame_height = image.height
266                    layers = 1
267
268            if frame_height * layers != image.height:
269                raise RuntimeError(
270                    f"Invalid texture size for variable #{name}:"
271                    + f" {frame_height}x{layers} != {image.height}"
272                    + f"\n  {info}"
273                )
274
275            logger.debug(f"Texture array: {image.width=}, {frame_height=}, {layers=}")
276            texture = self.ctx.texture_array(
277                size=(image.width, frame_height, layers),
278                components=4,
279                data=image.tobytes(),
280            )
281            texture.filter = (mgl.NEAREST, mgl.NEAREST)
282            texture.use(i)
283
284        # transform entire model
285
286        gui = model.display.get("gui") or DisplayPosition(
287            rotation=(30, 225, 0),
288            translation=(0, 0, 0),
289            scale=(0.625, 0.625, 0.625),
290        )
291
292        model_transform = cast(
293            Matrix44,
294            Matrix44.from_scale(gui.scale, "f4")
295            * get_rotation_matrix(gui.eulers)
296            * Matrix44.from_translation(gui.translation, "f4")
297            * Matrix44.from_translation((-8, -8, -8), "f4"),
298        )
299
300        normals_transform = Matrix44.from_y_rotation(-gui.eulers[1], "f4")
301        self.uniform("m_normals").write(normals_transform)
302
303        # render elements
304
305        baked_faces = list[BakedFace]()
306
307        for element in model.elements:
308            element_transform = model_transform.copy()
309
310            # TODO: rescale??
311            if rotation := element.rotation:
312                origin = np.array(rotation.origin)
313                element_transform *= cast(
314                    Matrix44,
315                    Matrix44.from_translation(origin, "f4")
316                    * get_rotation_matrix(rotation.eulers)
317                    * Matrix44.from_translation(-origin, "f4"),
318                )
319
320            # prepare each face of the element for rendering
321            for direction, face in element.faces.items():
322                baked_face = BakedFace(
323                    element=element,
324                    direction=direction,
325                    face=face,
326                    m_model=element_transform,
327                    texture0=texture_locs[face.texture_name],
328                    is_opaque=face.texture_name not in transparent_textures,
329                )
330                baked_faces.append(baked_face)
331
332        # TODO: use a map if this is actually slow
333        baked_faces.sort(key=lambda face: face.sortkey(self.eye))
334
335        for face in baked_faces:
336            self.uniform("m_model").write(face.m_model)
337            self.uniform("texture0").value = face.texture0
338
339            face.vao.render(self.face_prog)
340
341            if DebugType.NORMALS in debug:
342                self.uniform("m_model", self.debug_normal_prog).write(face.m_model)
343                face.vao.render(self.debug_normal_prog)
344
345        if DebugType.AXES in debug:
346            self.ctx.disable(mgl.CULL_FACE)
347            for axis, color in self.debug_axes:
348                self.uniform("color", self.debug_plane_prog).value = color
349                axis.render(self.debug_plane_prog)
350            self.ctx.enable(mgl.CULL_FACE)
351
352        self.ctx.finish()
353
354        # save to file
355
356        image = Image.frombytes(
357            mode="RGBA",
358            size=self.wnd.fbo.size,
359            data=self.wnd.fbo.read(components=4),
360        ).transpose(Image.Transpose.FLIP_TOP_BOTTOM)
361
362        output_path.parent.mkdir(parents=True, exist_ok=True)
363        image.save(output_path, format="png")
364
365    def uniform(self, name: str, program: Program | None = None):
366        program = program or self.face_prog
367        assert isinstance(uniform := program[name], Uniform)
368        return uniform
369
370
371@dataclass
372class BlockTextureInfo:
373    image_path: Path
374    meta: AnimationMeta | None
375
376
377@dataclass(kw_only=True)
378class BakedFace:
379    element: ModelElement
380    direction: FaceName
381    face: ElementFace
382    m_model: Matrix44
383    texture0: float
384    is_opaque: bool
385
386    def __post_init__(self):
387        self.verts = get_face_verts(self.element.from_, self.element.to, self.direction)
388
389        self.normals = get_face_normals(self.direction)
390
391        face_uv = self.face.uv or ElementFaceUV.default(self.element, self.direction)
392        self.uvs = [
393            value
394            for index in get_face_uv_indices(self.direction)
395            for value in face_uv.get_uv(index)
396        ]
397
398        self.vao = VAO()
399        self.vao.buffer(np.array(self.verts, np.float32), "3f", ["in_position"])
400        self.vao.buffer(np.array(self.normals, np.float32), "3f", ["in_normal"])
401        self.vao.buffer(np.array(self.uvs, np.float32) / 16, "2f", ["in_texcoord_0"])
402
403    @cached_property
404    def position(self):
405        x, y, z, n = 0, 0, 0, 0
406        for i in range(0, len(self.verts), 3):
407            x += self.verts[i]
408            y += self.verts[i + 1]
409            z += self.verts[i + 2]
410            n += 1
411        return (x / n, y / n, z / n)
412
413    def sortkey(self, eye: Vec3):
414        if self.is_opaque:
415            return 0
416        return sum((a - b) ** 2 for a, b in zip(eye, self.position))
417
418
419def get_face_verts(from_: Vec3, to: Vec3, direction: FaceName):
420    x1, y1, z1 = from_
421    x2, y2, z2 = to
422
423    # fmt: off
424    match direction:
425        case "south":
426            return [
427                x2, y1, z2,
428                x2, y2, z2,
429                x1, y1, z2,
430                x2, y2, z2,
431                x1, y2, z2,
432                x1, y1, z2,
433            ]
434        case "east":
435            return [
436                x2, y1, z1,
437                x2, y2, z1,
438                x2, y1, z2,
439                x2, y2, z1,
440                x2, y2, z2,
441                x2, y1, z2,
442            ]
443        case "down":
444            return [
445                x2, y1, z1,
446                x2, y1, z2,
447                x1, y1, z2,
448                x2, y1, z1,
449                x1, y1, z2,
450                x1, y1, z1,
451            ]
452        case "west":
453            return [
454                x1, y1, z2,
455                x1, y2, z2,
456                x1, y2, z1,
457                x1, y1, z2,
458                x1, y2, z1,
459                x1, y1, z1,
460            ]
461        case "north":
462            return [
463                x2, y2, z1,
464                x2, y1, z1,
465                x1, y1, z1,
466                x2, y2, z1,
467                x1, y1, z1,
468                x1, y2, z1,
469            ]
470        case "up":
471            return [
472                x2, y2, z1,
473                x1, y2, z1,
474                x2, y2, z2,
475                x1, y2, z1,
476                x1, y2, z2,
477                x2, y2, z2,
478            ]
479    # fmt: on
480
481
482def get_face_normals(direction: FaceName):
483    return 6 * get_direction_vec(direction)
484
485
486def get_face_uv_indices(direction: FaceName):
487    match direction:
488        case "south":
489            return (2, 3, 1, 3, 0, 1)
490        case "east":
491            return (2, 3, 1, 3, 0, 1)
492        case "down":
493            return (2, 3, 0, 2, 0, 1)
494        case "west":
495            return (2, 3, 0, 2, 0, 1)
496        case "north":
497            return (0, 1, 2, 0, 2, 3)
498        case "up":
499            return (3, 0, 2, 0, 1, 2)
500
501
502def orbit_camera(pitch: float, yaw: float):
503    """Both values are in degrees."""
504
505    eye = transform_vec(
506        (-64, 0, 0),
507        cast(
508            Matrix44,
509            Matrix44.identity(dtype="f4")
510            * Matrix44.from_y_rotation(math.radians(yaw))
511            * Matrix44.from_z_rotation(math.radians(pitch)),
512        ),
513    )
514
515    up = transform_vec(
516        (-1, 0, 0),
517        cast(
518            Matrix44,
519            Matrix44.identity(dtype="f4")
520            * Matrix44.from_y_rotation(math.radians(yaw))
521            * Matrix44.from_z_rotation(math.radians(90 - pitch)),
522        ),
523    )
524
525    return Matrix44.look_at(
526        eye=eye,
527        target=(0, 0, 0),
528        up=up,
529        dtype="f4",
530    ), eye
531
532
533def transform_vec(vec: Vec3, matrix: Matrix44) -> Vec3:
534    return np.matmul((*vec, 1), matrix, dtype="f4")[:3]
535
536
537def direction_camera(pos: FaceName, up: FaceName = "up"):
538    """eg. north -> camera is placed to the north of the model, looking south"""
539    eye = get_direction_vec(pos, 64)
540    return Matrix44.look_at(
541        eye=eye,
542        target=(0, 0, 0),
543        up=get_direction_vec(up),
544        dtype="f4",
545    ), eye
546
547
548def get_direction_vec(direction: FaceName, magnitude: float = 1):
549    match direction:
550        case "north":
551            return (0, 0, -magnitude)
552        case "south":
553            return (0, 0, magnitude)
554        case "west":
555            return (-magnitude, 0, 0)
556        case "east":
557            return (magnitude, 0, 0)
558        case "down":
559            return (0, -magnitude, 0)
560        case "up":
561            return (0, magnitude, 0)
562
563
564def read_shader(path: str, type: Literal["fragment", "vertex", "geometry"]):
565    file = resources.files(glsl) / path / f"{type}.glsl"
566    return file.read_text("utf-8")
567
568
569def get_rotation_matrix(eulers: Vec3) -> Matrix44:
570    return cast(
571        Matrix44,
572        Matrix44.from_x_rotation(-eulers[0], "f4")
573        * Matrix44.from_y_rotation(-eulers[1], "f4")
574        * Matrix44.from_z_rotation(-eulers[2], "f4"),
575    )