mirror of
https://github.com/3b1b/manim.git
synced 2025-04-13 09:47:07 +00:00
117 lines
3.2 KiB
Python
117 lines
3.2 KiB
Python
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)
|
|
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})"
|