mirror of
https://github.com/3b1b/manim.git
synced 2025-04-13 09:47:07 +00:00
95 lines
2.4 KiB
Python
95 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
from functools import lru_cache
|
|
import moderngl
|
|
from PIL import Image
|
|
|
|
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
|
|
|
|
|
|
ID_TO_TEXTURE: dict[int, moderngl.Texture] = 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(),
|
|
)
|
|
|
|
|
|
def get_texture_id(texture: moderngl.Texture) -> int:
|
|
tid = 0
|
|
while tid in ID_TO_TEXTURE:
|
|
tid += 1
|
|
ID_TO_TEXTURE[tid] = texture
|
|
texture.use(location=tid)
|
|
return tid
|
|
|
|
|
|
def release_texture(texture_id: int):
|
|
texture = ID_TO_TEXTURE.pop(texture_id, None)
|
|
if texture is not None:
|
|
texture.release()
|
|
|
|
|
|
@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,
|
|
)
|
|
|
|
|
|
@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})"
|