[Feature] Support Batch Textures (#2077)

* update batch texture

* update batch texture

* [BUG FIX] Fix URDF color overwrite. (#2065)

* Fix URDF color overwrite.
* Fix zero-copy being disabled for numpy.
* Add warning when enabling requires_grad while using GsTaichi dynamic arrays.
* More robust unit tests.

* [BUG FIX] Fix differentiable simulation support for ndarray. (#2068)

* Fix differentiable simulation support for ndarray.
* Fix shape issue with zero-copy for matrix field of shape Nx1.

* [MISC] Avoid performance penalty due to differentiability if disabled. (#2063)

* [MISC] Fix fastcache and zero-copy bugs. (#2050)

* Migrate to gstaichi 4.3.1.
* Enable zero-copy for ti.field except on Metal.
* Re-enable fastcache on CI.
* Get around memory initialization issue with field.
* Fix shape issue with zero-copy for matrix field of shape Nx1.
* Forcibly disable debug more when gradient is required.
* Add marker to set debug mode on each specific unit test.
* Fix missing copy.
* Fix collider state not reset at init.
* Fix unit test not supported on ARM runner.

---------

Co-authored-by: Alexis Duburcq <alexis.duburcq@gmail.com>

* [MISC] Bump Genesis version (v0.3.8). (#2073)

* [BUG FIX] Fix increased memory usage due to differentiable simulation. (#2074)

* Fix increased memory usage due to differentiable simulation.
* Cleanup differentiable rigid sim.

* [MISC] Speed up torch-based geom utils via 'torch.jit.script'. (#2075)

* More efficient GO2 env implementation.
* Speed up torch-based geom utils via 'torch.jit.script'.

* add unittest

* update unittest

* Update single_franka_batch_render.py

* Update pyproject.toml

* update snapshots

* Update image_exporter.py

* Delete utils_old.py

* update snapshots

* Update test_render.py

* debug

* update sorted textures

* debug(2)

* debug(3)

* debug(2)

* clear

* Update conftest.py

* update review

* update test_render

* Fallback to texture assets

---------

Co-authored-by: Alexis DUBURCQ <alexis.duburcq@gmail.com>
Co-authored-by: Sanghyun Son <shh1295@gmail.com>
Co-authored-by: Hugh Perkins <hughperkins@gmail.com>

Author: 4 people

examples/rigid/single_franka_batch_render.py

@@ -35,6 +35,7 @@ def main():
     ########################## entities ##########################
     plane = scene.add_entity(
         gs.morphs.Plane(),
+        surface=gs.surfaces.Default(diffuse_texture=gs.textures.BatchTexture.from_images(image_folder="textures")),
     )
     franka = scene.add_entity(
         gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),

genesis/options/surfaces.py

@@ -1,12 +1,13 @@
 from typing import Optional
+import math
 
 import numpy as np
 
 import genesis as gs
 
 from .misc import FoamOptions
 from .options import Options
-from .textures import ColorTexture, ImageTexture, Texture
+from .textures import Texture, ColorTexture, ImageTexture, BatchTexture
 
 
 ############################ Base ############################
@@ -215,52 +216,76 @@ def update_texture(
                 self.double_sided = double_sided
 
     def requires_uv(self):
-        return (
-            isinstance(self.get_texture(), ImageTexture)
-            or isinstance(self.opacity_texture, ImageTexture)
-            or isinstance(self.roughness_texture, ImageTexture)
-            or isinstance(self.metallic_texture, ImageTexture)
-            or isinstance(self.normal_texture, ImageTexture)
-            or isinstance(self.emissive_texture, ImageTexture)
+        textures = (
+            self.get_texture(),
+            self.opacity_texture,
+            self.roughness_texture,
+            self.metallic_texture,
+            self.normal_texture,
+            self.emissive_texture,
         )
+        return any(texture is not None and texture.requires_uv() for texture in textures)
+
+    def get_rgba(self, batch=False):
+        all_textures = []
+        for texture in (
+            self.get_texture() if self.emissive_texture is None else self.emissive_texture,
+            self.opacity_texture,
+        ):
+            textures = texture.textures if isinstance(texture, BatchTexture) else [texture]
+            all_textures.append(textures if batch else textures[:1])
+        color_textures, opacity_textures = all_textures
+
+        rgba_textures = []
+        num_colors = len(color_textures)
+        num_opacities = len(opacity_textures)
+        num_rgba = num_colors * num_opacities // math.gcd(num_colors, num_opacities)
+
+        for i in range(num_rgba):
+            color_texture = color_textures[i % num_colors]
+            opacity_texture = opacity_textures[i % num_opacities]
+
+            if isinstance(color_texture, ColorTexture):
+                if isinstance(opacity_texture, ColorTexture):
+                    rgba_texture = ColorTexture(color=(*color_texture.color, *opacity_texture.color))
+                elif isinstance(opacity_texture, ImageTexture) and opacity_texture.image_array is not None:
+                    rgb_color = np.round(np.array(color_texture.color) * 255).astype(np.uint8)
+                    rgb_array = np.full((*opacity_texture.image_array.shape[:2], 3), rgb_color, dtype=np.uint8)
+                    rgba_array = np.dstack((rgb_array, opacity_texture.image_array))
+                    rgba_scale = (1.0, 1.0, 1.0, *opacity_texture.image_color)
+                    rgba_texture = ImageTexture(image_array=rgba_array, image_color=rgba_scale)
+                else:
+                    rgba_texture = ColorTexture(color=(*color_texture.color, 1.0))
+
+            elif isinstance(color_texture, ImageTexture) and color_texture.image_array is not None:
+                if isinstance(opacity_texture, ColorTexture):
+                    a_color = np.round(np.array(opacity_texture.color) * 255).astype(np.uint8)
+                    a_array = np.full((*color_texture.image_array.shape[:2],), a_color, dtype=np.uint8)
+                    rgba_array = np.dstack((color_texture.image_array, a_array))
+                    rgba_scale = (*color_texture.image_color, 1.0)
+                elif (
+                    isinstance(opacity_texture, ImageTexture)
+                    and opacity_texture.image_array is not None
+                    and opacity_texture.image_array.shape[:2] == color_texture.image_array.shape[:2]
+                ):
+                    rgba_array = np.dstack((color_texture.image_array, opacity_texture.image_array))
+                    rgba_scale = (*color_texture.image_color, *opacity_texture.image_color)
+                else:
+                    if isinstance(opacity_texture, ImageTexture) and opacity_texture.image_array is not None:
+                        gs.logger.warning(
+                            "Color and opacity image shapes do not match. Fall back to fully opaque texture."
+                        )
+                    a_array = np.full(color_texture.image_array.shape[:2], 255, dtype=np.uint8)
+                    rgba_array = np.dstack((color_texture.image_array, a_array))
+                    rgba_scale = (*color_texture.image_color, 1.0)
+                rgba_texture = ImageTexture(image_array=rgba_array, image_color=rgba_scale)
 
-    def get_rgba(self):
-        texture = self.get_texture() if self.emissive_texture is None else self.emissive_texture
-        opacity_texture = self.opacity_texture
-
-        if isinstance(texture, ColorTexture):
-            if isinstance(opacity_texture, ColorTexture):
-                return ColorTexture(color=(*texture.color, *opacity_texture.color))
-            if isinstance(opacity_texture, ImageTexture) and opacity_texture.image_array is not None:
-                rgb_color = np.round(np.array(texture.color) * 255).astype(np.uint8)
-                rgb_array = np.full((*opacity_texture.image_array.shape[:2], 3), rgb_color, dtype=np.uint8)
-                rgba_array = np.dstack((rgb_array, opacity_texture.image_array))
-                rgba_scale = (1.0, 1.0, 1.0, *opacity_texture.image_color)
-                return ImageTexture(image_array=rgba_array, image_color=rgba_scale)
-            return ColorTexture(color=(*texture.color, 1.0))
-
-        if isinstance(texture, ImageTexture) and texture.image_array is not None:
-            if isinstance(opacity_texture, ColorTexture):
-                a_color = np.round(np.array(opacity_texture.color) * 255).astype(np.uint8)
-                a_array = np.full((*texture.image_array.shape[:2],), a_color, dtype=np.uint8)
-                rgba_array = np.dstack((texture.image_array, a_array))
-                rgba_scale = (*texture.image_color, 1.0)
-            elif (
-                isinstance(opacity_texture, ImageTexture)
-                and opacity_texture.image_array is not None
-                and opacity_texture.image_array.shape[:2] == texture.image_array.shape[:2]
-            ):
-                rgba_array = np.dstack((texture.image_array, opacity_texture.image_array))
-                rgba_scale = (*texture.image_color, *opacity_texture.image_color)
             else:
-                if isinstance(opacity_texture, ImageTexture) and opacity_texture.image_array is not None:
-                    gs.logger.warning("Color and opacity image shapes do not match. Fall back to fully opaque texture.")
-                a_array = np.full(texture.image_array.shape[:2], 255, dtype=np.uint8)
-                rgba_array = np.dstack((texture.image_array, a_array))
-                rgba_scale = (*texture.image_color, 1.0)
-            return ImageTexture(image_array=rgba_array, image_color=rgba_scale)
-
-        return ColorTexture(color=(1.0, 1.0, 1.0, 1.0))
+                rgba_texture = ColorTexture(color=(1.0, 1.0, 1.0, 1.0))
+
+            rgba_textures.append(rgba_texture)
+
+        return BatchTexture(textures=rgba_textures) if batch else rgba_textures[0]
 
     def set_texture(self, texture):
         raise NotImplementedError

genesis/options/textures.py

@@ -1,16 +1,17 @@
 import os
 from typing import Optional, List, Union
 
-import Imath
 import numpy as np
-import OpenEXR
 from PIL import Image
 
 import genesis as gs
 import genesis.utils.mesh as mu
 
 from .options import Options
 
+IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".bmp", ".webp", ".hdr", ".exr")
+HDR_EXTENSIONS = (".hdr", ".exr")
+
 
 class Texture(Options):
     """
@@ -36,6 +37,9 @@ def apply_cutoff(self, cutoff):
     def is_black(self):
         raise NotImplementedError
 
+    def requires_uv(self):
+        raise NotImplementedError
+
 
 class ColorTexture(Texture):
     """
@@ -74,6 +78,9 @@ def apply_cutoff(self, cutoff):
     def is_black(self):
         return all(c < gs.EPS for c in self.color)
 
+    def requires_uv(self):
+        return False
+
 
 class ImageTexture(Texture):
     """
@@ -119,7 +126,7 @@ def __init__(self, **data):
                 )
 
             # Load image_path as actual image_array, unless for special texture images (e.g. `.hdr` and `.exr`) that are only supported by raytracers
-            if self.image_path.endswith((".hdr", ".exr")):
+            if self.image_path.endswith(HDR_EXTENSIONS):
                 self.encoding = "linear"  # .exr or .hdr images should be encoded with 'linear'
                 if self.image_path.endswith((".exr")):
                     self.image_path = mu.check_exr_compression(self.image_path)
@@ -186,16 +193,139 @@ def check_simplify(self):
         else:
             return self
 
+    def apply_cutoff(self, cutoff):
+        if cutoff is None or self.image_array is None:  # Cutoff does not apply on image file.
+            return
+        self.image_array = np.where(self.image_array >= 255.0 * cutoff, 255, 0).astype(np.uint8)
+
+    def is_black(self):
+        return all(c < gs.EPS for c in self.image_color) or np.max(self.image_array) == 0
+
+    def requires_uv(self):
+        return True
+
     def mean_color(self):
         return self._mean_color
 
     def channel(self):
         return self._channel
 
+
+class BatchTexture(Texture):
+    """
+    A batch of textures for batch rendering.
+
+    Parameters
+    ----------
+    textures : List[Optional[Texture]]
+        List of textures.
+    """
+
+    textures: Optional[List[Optional[Texture]]]
+
+    def __init__(self, **data):
+        super().__init__(**data)
+
+    @staticmethod
+    def from_images(
+        image_paths: Optional[List[str]] = None,
+        image_folder: Optional[str] = None,
+        image_arrays: Optional[List[np.ndarray]] = None,
+        image_colors: Optional[Union[List[float], List[List[float]]]] = None,
+        encoding: str = "srgb",
+    ):
+        """
+        Create a batch texture from images.
+
+        Parameters
+        ----------
+        image_paths : List[str], optional
+            List of paths to the image files.
+        image_folder : str, optional
+            Path to the image folder.
+        image_arrays : List[np.ndarray], optional
+            List of image arrays.
+        image_colors : List[Union[float, List[float]]], optional
+            List of color factors that will be multiplied with the base color, stored as tuple. Default is None.
+        encoding : str, optional
+            The encoding way of the image. Possible values are ['srgb', 'linear']. Default is 'srgb'.
+
+            - 'srgb': Encoding of some RGB images.
+            - 'linear': All generic images, such as opacity, roughness and normal, should be encoded with 'linear'.
+        """
+        image_sources = (image_paths, image_folder, image_arrays)
+        if sum(x is not None for x in image_sources) != 1:
+            gs.raise_exception("Please set exactly one of `image_paths`, `image_folder`, `image_arrays`.")
+
+        image_textures = []
+        if image_folder is not None:
+            input_image_folder = image_folder
+            if not os.path.exists(image_folder):
+                image_folder = os.path.join(gs.utils.get_assets_dir(), image_folder)
+            if not os.path.exists(image_folder):
+                gs.raise_exception(
+                    f"Directory not found in either current directory or assets directory: '{input_image_folder}'."
+                )
+            image_paths = [
+                os.path.join(image_folder, image_path)
+                for image_path in sorted(os.listdir(image_folder))
+                if image_path.lower().endswith(IMAGE_EXTENSIONS)
+            ]
+
+        num_images = len(image_paths) if image_paths is not None else len(image_arrays)
+        if num_images == 0:
+            gs.raise_exception("No images found.")
+
+        if image_colors is not None:
+            if isinstance(image_colors[0], float):  # List[float]
+                image_colors = [image_colors for _ in range(num_images)]
+            else:  # List[List[float]]
+                if len(image_colors) != num_images:
+                    gs.raise_exception("The number of image colors must be the same as the number of images.")
+        else:
+            image_colors = [None] * num_images
+
+        if image_paths is not None:
+            for image_path, image_color in zip(image_paths, image_colors):
+                image_textures.append(ImageTexture(image_path=image_path, image_color=image_color, encoding=encoding))
+        else:
+            for image_array, image_color in zip(image_arrays, image_colors):
+                image_textures.append(ImageTexture(image_array=image_array, image_color=image_color, encoding=encoding))
+
+        return BatchTexture(textures=image_textures)
+
+    @staticmethod
+    def from_colors(
+        colors: List[List[float]],
+    ):
+        """
+        Create a batch texture from colors.
+
+        Parameters
+        ----------
+        colors : List[List[float]]
+            List of colors.
+        """
+        color_textures = [ColorTexture(color=color) for color in colors]
+        return BatchTexture(textures=color_textures)
+
+    def check_dim(self, dim):
+        return BatchTexture(textures=[texture.check_dim(dim) for texture in self.textures]).check_simplify()
+
+    def check_simplify(self):
+        self.textures = [texture.check_simplify() if texture is not None else None for texture in self.textures]
+        return self
+
     def apply_cutoff(self, cutoff):
-        if cutoff is None or self.image_array is None:  # Cutoff does not apply on image file.
-            return
-        self.image_array = np.where(self.image_array >= 255.0 * cutoff, 255, 0).astype(np.uint8)
+        for texture in self.textures:
+            if texture is not None:
+                texture.apply_cutoff(cutoff)
 
     def is_black(self):
-        return all(c < gs.EPS for c in self.image_color) or np.max(self.image_array) == 0
+        return all(texture.is_black() if texture is not None else True for texture in self.textures)
+
+    def requires_uv(self):
+        return any(texture is not None and texture.requires_uv() for texture in self.textures)
+
+    def merge(self, other: "BatchTexture"):
+        self.textures.extend(other.textures)