Implement OpenGLImageMobject (#2534)

* initial commit

make OpenGLImageMobject a subclass of OpenGLTexturedSurface

* added typing

* unused import

* explicit imports

* new typing style

* convert to correct type

* update file structure

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Added support for Path and image_mode parameter

* remove duplicate function

* let moderngl accept PIL Images only

* allow custom resampling

* allow greyscale

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Naveen M K <[email protected]>

Author: 3 people

manim/mobject/opengl/opengl_image_mobject.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+__all__ = [
+    "OpenGLImageMobject",
+]
+
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+from manim.mobject.opengl.opengl_surface import OpenGLSurface, OpenGLTexturedSurface
+from manim.utils.images import get_full_raster_image_path
+
+
+class OpenGLImageMobject(OpenGLTexturedSurface):
+    def __init__(
+        self,
+        filename_or_array: str | Path | np.ndarray,
+        width: float = None,
+        height: float = None,
+        image_mode: str = "RGBA",
+        resampling_algorithm: int = Image.BICUBIC,
+        opacity: float = 1,
+        gloss: float = 0,
+        shadow: float = 0,
+        **kwargs,
+    ):
+        self.image = filename_or_array
+        self.resampling_algorithm = resampling_algorithm
+        if type(filename_or_array) == np.ndarray:
+            self.size = self.image.shape[1::-1]
+        elif isinstance(filename_or_array, (str, Path)):
+            path = get_full_raster_image_path(filename_or_array)
+            self.size = Image.open(path).size
+
+        if width is None and height is None:
+            width = 4 * self.size[0] / self.size[1]
+            height = 4
+        if height is None:
+            height = width * self.size[1] / self.size[0]
+        if width is None:
+            width = height * self.size[0] / self.size[1]
+
+        surface = OpenGLSurface(
+            lambda u, v: np.array([u, v, 0]),
+            [-width / 2, width / 2],
+            [-height / 2, height / 2],
+            opacity=opacity,
+            gloss=gloss,
+            shadow=shadow,
+        )
+
+        super().__init__(
+            surface,
+            self.image,
+            image_mode=image_mode,
+            opacity=opacity,
+            gloss=gloss,
+            shadow=shadow,
+            **kwargs,
+        )
+
+    def get_image_from_file(
+        self,
+        image_file: str | Path | np.ndarray,
+        image_mode: str,
+    ):
+        if isinstance(image_file, (str, Path)):
+            return super().get_image_from_file(image_file, image_mode)
+        else:
+            return (
+                Image.fromarray(image_file.astype("uint8"))
+                .convert(image_mode)
+                .resize(
+                    np.array(image_file.shape[:2])
+                    * 200,  # assumption of 200 ppmu (pixels per manim unit) would suffice
+                    resample=self.resampling_algorithm,
+                )
+            )

manim/mobject/opengl/opengl_surface.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
+from pathlib import Path
+from typing import Iterable
+
 import moderngl
 import numpy as np
+from PIL import Image
 
 from manim.constants import *
 from manim.mobject.opengl.opengl_mobject import OpenGLMobject
 from manim.utils.bezier import integer_interpolate, interpolate
 from manim.utils.color import *
 from manim.utils.config_ops import _Data, _Uniforms
-from manim.utils.images import get_full_raster_image_path
+from manim.utils.images import change_to_rgba_array, get_full_raster_image_path
 from manim.utils.iterables import listify
 from manim.utils.space_ops import normalize_along_axis
 
@@ -324,22 +328,37 @@ class OpenGLTexturedSurface(OpenGLSurface):
     num_textures = _Uniforms()
 
     def __init__(
-        self, uv_surface, image_file, dark_image_file=None, shader_folder=None, **kwargs
+        self,
+        uv_surface: OpenGLSurface,
+        image_file: str | Path,
+        dark_image_file: str | Path = None,
+        image_mode: str | Iterable[str] = "RGBA",
+        shader_folder: str | Path = None,
+        **kwargs,
     ):
         self.uniforms = {}
 
         if not isinstance(uv_surface, OpenGLSurface):
             raise Exception("uv_surface must be of type OpenGLSurface")
+        if type(image_file) == np.ndarray:
+            image_file = change_to_rgba_array(image_file)
+
         # Set texture information
-        if dark_image_file is None:
-            dark_image_file = image_file
-            self.num_textures = 1
-        else:
-            self.num_textures = 2
+        if isinstance(image_mode, (str, Path)):
+            image_mode = [image_mode] * 2
+        image_mode_light, image_mode_dark = image_mode
         texture_paths = {
-            "LightTexture": get_full_raster_image_path(image_file),
-            "DarkTexture": get_full_raster_image_path(dark_image_file),
+            "LightTexture": self.get_image_from_file(
+                image_file,
+                image_mode_light,
+            ),
+            "DarkTexture": self.get_image_from_file(
+                dark_image_file or image_file,
+                image_mode_dark,
+            ),
         }
+        if dark_image_file:
+            self.num_textures = 2
 
         self.uv_surface = uv_surface
         self.uv_func = uv_surface.uv_func
@@ -349,6 +368,14 @@ def __init__(
         self.gloss = self.uv_surface.gloss
         super().__init__(texture_paths=texture_paths, **kwargs)
 
+    def get_image_from_file(
+        self,
+        image_file: str | Path,
+        image_mode: str,
+    ):
+        image_file = get_full_raster_image_path(image_file)
+        return Image.open(image_file).convert(image_mode)
+
     def init_data(self):
         super().init_data()
         self.im_coords = np.zeros((0, 2))

manim/mobject/types/image_mobject.py

@@ -17,7 +17,7 @@
 from ...mobject.mobject import Mobject
 from ...utils.bezier import interpolate
 from ...utils.color import WHITE, color_to_int_rgb
-from ...utils.images import get_full_raster_image_path
+from ...utils.images import change_to_rgba_array, get_full_raster_image_path
 
 
 class AbstractImageMobject(Mobject):
@@ -183,26 +183,13 @@ def __init__(
         else:
             self.pixel_array = np.array(filename_or_array)
         self.pixel_array_dtype = kwargs.get("pixel_array_dtype", "uint8")
-        self.change_to_rgba_array()
+        self.pixel_array = change_to_rgba_array(
+            self.pixel_array, self.pixel_array_dtype
+        )
         if self.invert:
             self.pixel_array[:, :, :3] = 255 - self.pixel_array[:, :, :3]
         super().__init__(scale_to_resolution, **kwargs)
 
-    def change_to_rgba_array(self):
-        """Converts an RGB array into RGBA with the alpha value opacity maxed."""
-        pa = self.pixel_array
-        if len(pa.shape) == 2:
-            pa = pa.reshape(list(pa.shape) + [1])
-        if pa.shape[2] == 1:
-            pa = pa.repeat(3, axis=2)
-        if pa.shape[2] == 3:
-            alphas = 255 * np.ones(
-                list(pa.shape[:2]) + [1],
-                dtype=self.pixel_array_dtype,
-            )
-            pa = np.append(pa, alphas, axis=2)
-        self.pixel_array = pa
-
     def get_pixel_array(self):
         """A simple getter method."""
         return self.pixel_array

manim/opengl/__init__.py

@@ -7,6 +7,7 @@
 
 
 from manim.mobject.opengl.dot_cloud import *
+from manim.mobject.opengl.opengl_image_mobject import *
 from manim.mobject.opengl.opengl_mobject import *
 from manim.mobject.opengl.opengl_point_cloud_mobject import *
 from manim.mobject.opengl.opengl_surface import *

manim/renderer/opengl_renderer.py

@@ -381,18 +381,21 @@ def render_mobject(self, mobject):
             mesh.render()
 
     def get_texture_id(self, path):
-        if path not in self.path_to_texture_id:
-            # A way to increase tid's sequentially
+        if repr(path) not in self.path_to_texture_id:
             tid = len(self.path_to_texture_id)
-            im = Image.open(path)
             texture = self.context.texture(
-                size=im.size,
-                components=len(im.getbands()),
-                data=im.tobytes(),
+                size=path.size,
+                components=len(path.getbands()),
+                data=path.tobytes(),
             )
+            texture.repeat_x = False
+            texture.repeat_y = False
+            texture.filter = (moderngl.NEAREST, moderngl.NEAREST)
+            texture.swizzle = "RRR1" if path.mode == "L" else "RGBA"
             texture.use(location=tid)
-            self.path_to_texture_id[path] = tid
-        return self.path_to_texture_id[path]
+            self.path_to_texture_id[repr(path)] = tid
+
+        return self.path_to_texture_id[repr(path)]
 
     def update_skipping_status(self):
         """

manim/utils/images.py

@@ -2,7 +2,12 @@
 
 from __future__ import annotations
 
-__all__ = ["get_full_raster_image_path", "drag_pixels", "invert_image"]
+__all__ = [
+    "get_full_raster_image_path",
+    "drag_pixels",
+    "invert_image",
+    "change_to_rgba_array",
+]
 
 import numpy as np
 from PIL import Image
@@ -32,3 +37,19 @@ def invert_image(image: np.array) -> Image:
     arr = np.array(image)
     arr = (255 * np.ones(arr.shape)).astype(arr.dtype) - arr
     return Image.fromarray(arr)
+
+
+def change_to_rgba_array(image, dtype="uint8"):
+    """Converts an RGB array into RGBA with the alpha value opacity maxed."""
+    pa = image
+    if len(pa.shape) == 2:
+        pa = pa.reshape(list(pa.shape) + [1])
+    if pa.shape[2] == 1:
+        pa = pa.repeat(3, axis=2)
+    if pa.shape[2] == 3:
+        alphas = 255 * np.ones(
+            list(pa.shape[:2]) + [1],
+            dtype=dtype,
+        )
+        pa = np.append(pa, alphas, axis=2)
+    return pa