123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117 |
- from __future__ import annotations
- import os
- import re
- from functools import lru_cache
- import moderngl
- from PIL import Image
- import numpy as np
- from manimlib.utils.directories import get_shader_dir
- from manimlib.utils.file_ops import find_file
- from typing import TYPE_CHECKING
- if TYPE_CHECKING:
- from typing import Sequence, Optional, Tuple
- from manimlib.typing import UniformDict
- from moderngl.vertex_array import VertexArray
- from moderngl.framebuffer import Framebuffer
- # Global maps to reflect uniform status
- PROGRAM_UNIFORM_MIRRORS: dict[int, dict[str, float | tuple]] = dict()
- @lru_cache()
- def image_path_to_texture(path: str, ctx: moderngl.Context) -> moderngl.Texture:
- im = Image.open(path).convert("RGBA")
- return ctx.texture(
- size=im.size,
- components=len(im.getbands()),
- data=im.tobytes(),
- )
- @lru_cache()
- def get_shader_program(
- ctx: moderngl.context.Context,
- vertex_shader: str,
- fragment_shader: Optional[str] = None,
- geometry_shader: Optional[str] = None,
- ) -> moderngl.Program:
- return ctx.program(
- vertex_shader=vertex_shader,
- fragment_shader=fragment_shader,
- geometry_shader=geometry_shader,
- )
- def set_program_uniform(
- program: moderngl.Program,
- name: str,
- value: float | tuple | np.ndarray
- ) -> bool:
- """
- Sets a program uniform, and also keeps track of a dictionary
- of previously set uniforms for that program so that it
- doesn't needlessly reset it, requiring an exchange with gpu
- memory, if it sees the same value again.
- Returns True if changed the program, False if it left it as is.
- """
- pid = id(program)
- if pid not in PROGRAM_UNIFORM_MIRRORS:
- PROGRAM_UNIFORM_MIRRORS[pid] = dict()
- uniform_mirror = PROGRAM_UNIFORM_MIRRORS[pid]
- if type(value) is np.ndarray and value.ndim > 0:
- value = tuple(value.flatten())
- if uniform_mirror.get(name, None) == value:
- return False
- try:
- program[name].value = value
- except KeyError:
- return False
- uniform_mirror[name] = value
- return True
- @lru_cache()
- def get_shader_code_from_file(filename: str) -> str | None:
- if not filename:
- return None
- try:
- filepath = find_file(
- filename,
- directories=[get_shader_dir(), "/"],
- extensions=[],
- )
- except IOError:
- return None
- with open(filepath, "r") as f:
- result = f.read()
- # To share functionality between shaders, some functions are read in
- # from other files an inserted into the relevant strings before
- # passing to ctx.program for compiling
- # Replace "#INSERT " lines with relevant code
- insertions = re.findall(r"^#INSERT .*\.glsl$", result, flags=re.MULTILINE)
- for line in insertions:
- inserted_code = get_shader_code_from_file(
- os.path.join("inserts", line.replace("#INSERT ", ""))
- )
- result = result.replace(line, inserted_code)
- return result
- def get_colormap_code(rgb_list: Sequence[float]) -> str:
- data = ",".join(
- "vec3({}, {}, {})".format(*rgb)
- for rgb in rgb_list
- )
- return f"vec3[{len(rgb_list)}]({data})"
|