diff --git a/.gitignore b/.gitignore index 9b09b257..0866e138 100644 --- a/.gitignore +++ b/.gitignore @@ -146,4 +146,9 @@ wfc_binding.cpp .gen_files/ .config/ + +# Mac +*.icloud + +# Tests thom test.gltf diff --git a/examples/procgen_grid.py b/examples/procgen_grid.py index 3a3e42ce..355ea3c5 100644 --- a/examples/procgen_grid.py +++ b/examples/procgen_grid.py @@ -18,20 +18,36 @@ ) * 0.6 ) -scene += sm.ProcgenGrid(specific_map=specific_map) + +proc_grid = sm.ProcgenGrid(specific_map=specific_map) + +# Example using a predefined colormap from matplotlib +# See https://matplotlib.org/stable/api/cm_api.html#matplotlib.cm.get_cmap +proc_grid.add_texture_cmap_along_axis(axis="y", cmap="viridis", n_colors=5) +scene += proc_grid + scene += sm.LightSun() -scene.show() +scene.show(show_edges=True) -input("Press Enter for second scene") +# input("Press Enter for second scene") scene.close() scene.clear() # Second scene: generating from this map scene += sm.ProcgenGrid(width=3, height=3, sample_map=specific_map) + +# Example creating our own colormap +# https://matplotlib.org/stable/tutorials/colors/colormap-manipulation.html +from matplotlib.colors import ListedColormap + + +cmap = ListedColormap(["red", "blue", "green"]) +scene.tree_children[0].add_texture_cmap_along_axis(axis="x", cmap=cmap) + scene += sm.LightSun() scene.show() -input("Press Enter for third scene") +# input("Press Enter for third scene") scene.close() scene.clear() @@ -47,8 +63,9 @@ neighbors = [(tiles[1], tiles[0]), (tiles[0], tiles[0]), (tiles[1], tiles[1])] scene += sm.ProcgenGrid(width=3, height=3, tiles=tiles, neighbors=neighbors, weights=weights, symmetries=symmetries) scene += sm.LightSun() +scene.tree_children[0].add_texture_cmap_along_axis(axis="x", cmap="viridis") scene.show() -input("Press Enter to close") +# input("Press Enter to close") scene.close() diff --git a/src/simenv/assets/__init__.py b/src/simenv/assets/__init__.py index 9d35c162..3044f36d 100644 --- a/src/simenv/assets/__init__.py +++ b/src/simenv/assets/__init__.py @@ -2,6 +2,7 @@ from .asset import Asset from .camera import Camera, CameraDistant from .collider import * +from .colors import * from .light import * from .material import * from .object import * diff --git a/src/simenv/assets/colors.py b/src/simenv/assets/colors.py new file mode 100644 index 00000000..3c9c7a1a --- /dev/null +++ b/src/simenv/assets/colors.py @@ -0,0 +1,88 @@ +# Copyright 2022 The HuggingFace Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Lint as: python3 +""" Some predefined RGB colors.""" +import numpy as np + + +BLACK = (0.0, 0.0, 0.0) +BLUE = (0.0, 0.0, 1.0) +CYAN = (0.0, 1.0, 1.0) +GRAY25 = (0.25, 0.25, 0.25) +GRAY50 = (0.5, 0.5, 0.5) +GRAY75 = (0.75, 0.75, 0.75) +GRAY = GRAY50 +GREEN = (0.0, 1.0, 0.0) +MAGENTA = (1.0, 0.0, 1.0) +OLIVE = (0.5, 0.5, 0.0) +PURPLE = (0.5, 0.0, 0.5) +RED = (1.0, 0.0, 0.0) +TEAL = (0.0, 0.5, 0.5) +WHITE = (1.0, 1.0, 1.0) +YELLOW = (1.0, 1.0, 0.0) + +COLORS_ALL = { + "BLACK": BLACK, + "BLUE": BLUE, + "CYAN": CYAN, + "GRAY25": GRAY25, + "GRAY50": GRAY50, + "GRAY75": GRAY75, + "GRAY": GRAY, + "GREEN": GREEN, + "MAGENTA": MAGENTA, + "OLIVE": OLIVE, + "PURPLE": PURPLE, + "RED": RED, + "TEAL": TEAL, + "WHITE": WHITE, + "YELLOW": YELLOW, +} + +COLORS_NO_GRAYSCALE = { + "BLUE": BLUE, + "CYAN": CYAN, + "GREEN": GREEN, + "MAGENTA": MAGENTA, + "OLIVE": OLIVE, + "PURPLE": PURPLE, + "RED": RED, + "TEAL": TEAL, + "YELLOW": YELLOW, +} + +COLORS_ONLY_GRAYSCALE = { + "WHITE": WHITE, + "GRAY25": GRAY25, + "GRAY50": GRAY50, + "GRAY75": GRAY75, + "GRAY100": BLACK, +} + +CMAP_ALL = np.array([list(COLORS_ALL.values())] * 2) # Final shape: (2, len(colors), 3) for (U, V, RGB) +CMAP_3_COLORS = np.array([[GREEN, GREEN]] * 2) # Final shape: (2, len(colors), 3) for (U, V, RGB) +CMAP_ONLY_COLORS = np.array( + [list(COLORS_NO_GRAYSCALE.values())] * 2 +) # Final shape: (2, len(colors), 3) for (U, V, RGB) +CMAP_ONLY_GRAYSCALE = np.array( + [list(COLORS_ONLY_GRAYSCALE.values())] * 2 +) # Final shape: (2, len(colors), 3) for (U, V, RGB) + +TEXTURES_ALL = { + "cmap_all": CMAP_ALL, + "cmap_3_colors": CMAP_3_COLORS, + "cmap_only_colors": CMAP_ONLY_COLORS, + "cmap_only_grayscale": CMAP_ONLY_GRAYSCALE, +} diff --git a/src/simenv/assets/material.py b/src/simenv/assets/material.py index 90818ae7..62c7b94a 100644 --- a/src/simenv/assets/material.py +++ b/src/simenv/assets/material.py @@ -17,22 +17,13 @@ import copy import itertools from dataclasses import dataclass -from typing import ClassVar, List, Optional +from typing import ClassVar, List, Optional, Union import numpy as np import pyvista -from .utils import camelcase_to_snakecase - - -class classproperty(object): - # required to use a classmethod as a property - # see https://stackoverflow.com/questions/128573/using-property-on-classmethods - def __init__(self, fget): - self.fget = fget - - def __get__(self, owner_self, owner_cls): - return self.fget(owner_cls) +from . import colors +from .utils import camelcase_to_snakecase, classproperty # TODO thom this is a very basic PBR Metrial class, mostly here to be able to load a gltf - strongly base on GLTF definitions @@ -101,14 +92,14 @@ class Material: __NEW_ID: ClassVar[int] = itertools.count() # Singleton to count instances of the classes for automatic naming base_color: Optional[List[float]] = None - base_color_texture: Optional[pyvista.Texture] = None + base_color_texture: Optional[Union[np.ndarray, pyvista.Texture]] = None metallic_factor: Optional[float] = None roughness_factor: Optional[float] = None - metallic_roughness_texture: Optional[pyvista.Texture] = None + metallic_roughness_texture: Optional[Union[np.ndarray, pyvista.Texture]] = None - normal_texture: Optional[pyvista.Texture] = None - occlusion_texture: Optional[pyvista.Texture] = None - emissive_texture: Optional[pyvista.Texture] = None + normal_texture: Optional[Union[np.ndarray, pyvista.Texture]] = None + occlusion_texture: Optional[Union[np.ndarray, pyvista.Texture]] = None + emissive_texture: Optional[Union[np.ndarray, pyvista.Texture]] = None emissive_factor: Optional[List[float]] = None alpha_mode: Optional[str] = None alpha_cutoff: Optional[float] = None @@ -118,6 +109,21 @@ class Material: def __post_init__(self): # Setup all our default values + + # Convert numpy array textures to + for tex in [ + "base_color_texture", + "metallic_roughness_texture", + "normal_texture", + "occlusion_texture", + "emissive_texture", + ]: + tex_value = getattr(self, tex, None) + if isinstance(tex_value, np.ndarray): + if tex_value.ndim != 3: + raise ValueError(f"{tex} must be a 3D numpy array (U, V, RGB)") + setattr(self, tex, pyvista.Texture(tex_value)) + if self.base_color is None: self.base_color = [1.0, 1.0, 1.0, 1.0] elif isinstance(self.base_color, np.ndarray): @@ -168,60 +174,64 @@ def copy(self): # Various default colors @classproperty def RED(cls): - return cls(base_color=(1.0, 0.0, 0.0)) + return cls(base_color=colors.RED) @classproperty def GREEN(cls): - return cls(base_color=(0.0, 1.0, 0.0)) + return cls(base_color=colors.GREEN) @classproperty def BLUE(cls): - return cls(base_color=(0.0, 0.0, 1.0)) + return cls(base_color=colors.BLUE) @classproperty def CYAN(cls): - return cls(base_color=(0.0, 1.0, 1.0)) + return cls(base_color=colors.CYAN) @classproperty def MAGENTA(cls): - return cls(base_color=(1.0, 0.0, 1.0)) + return cls(base_color=colors.MAGENTA) @classproperty def YELLOW(cls): - return cls(base_color=(1.0, 1.0, 0.0)) + return cls(base_color=colors.YELLOW) @classproperty def BLACK(cls): - return cls(base_color=(0.0, 0.0, 0.0)) + return cls(base_color=colors.BLACK) @classproperty def WHITE(cls): - return cls(base_color=(1.0, 1.0, 1.0)) + return cls(base_color=colors.WHITE) @classproperty def GRAY(cls): - return cls.GRAY50 + return cls(base_color=colors.GRAY) @classproperty def GRAY25(cls): - return cls(base_color=(0.25, 0.25, 0.25)) + return cls(base_color=colors.GRAY25) @classproperty def GRAY50(cls): - return cls(base_color=(0.5, 0.5, 0.5)) + return cls(base_color=colors.GRAY50) @classproperty def GRAY75(cls): - return cls(base_color=(0.75, 0.75, 0.75)) + return cls(base_color=colors.GRAY75) @classproperty def TEAL(cls): - return cls(base_color=(0.0, 0.5, 0.5)) + return cls(base_color=colors.TEAL) @classproperty def PURPLE(cls): - return cls(base_color=(0.5, 0.0, 0.5)) + return cls(base_color=colors.PURPLE) @classproperty def OLIVE(cls): - return cls(base_color=(0.5, 0.5, 0.0)) + return cls(base_color=colors.OLIVE) + + @classproperty + def CMAP_3_COLORS(cls): + return cls(base_color_texture=colors.CMAP_3_COLORS) diff --git a/src/simenv/assets/object.py b/src/simenv/assets/object.py index 1c672b8f..59b58c1d 100644 --- a/src/simenv/assets/object.py +++ b/src/simenv/assets/object.py @@ -19,6 +19,8 @@ import numpy as np import pyvista as pv +from matplotlib.cm import get_cmap +from matplotlib.colors import Colormap from .asset import Asset from .collider import Collider @@ -45,7 +47,7 @@ class Object3D(Asset): def __init__( self, - mesh: Optional[pv.UnstructuredGrid] = None, + mesh: Optional[Union[pv.UnstructuredGrid, pv.PolyData]] = None, material: Optional[Material] = None, name: Optional[str] = None, position: Optional[List[float]] = None, @@ -56,14 +58,40 @@ def __init__( ): super().__init__(name=name, position=position, parent=parent, children=children, collider=collider, **kwargs) - self.mesh = mesh if mesh is not None else pv.PolyData() + self._mesh = None + self.mesh = mesh # Avoid having averaging normals at shared points # (default pyvista behavior:https://docs.pyvista.org/api/core/_autosummary/pyvista.PolyData.compute_normals.html) if self.mesh is not None: self.mesh.compute_normals(inplace=True, cell_normals=False, split_vertices=True) - self.material = material if material is not None else Material() + self._material = None + self.material = material + + @property + def material(self): + return self._material + + @material.setter + def material(self, material: Material): + if material is None: + material = Material() + if not isinstance(material, Material): + raise ValueError(f"{material} is not a sm.Material instance") + self._material = material + + @property + def mesh(self): + return self._mesh + + @mesh.setter + def mesh(self, mesh: pv.PolyData): + if mesh is None: + mesh = pv.PolyData() + if not isinstance(mesh, (pv.UnstructuredGrid, pv.PolyData)): + raise ValueError(f"{mesh} is not a pyvista.PolyData or pyvista.UnstructuredGrid instance") + self._mesh = mesh def copy(self, with_children=True, **kwargs): """Copy an Object3D node in a new (returned) object. @@ -1053,11 +1081,65 @@ def __init__( z = np.array(z) # If it is a structured grid, extract the surface mesh (PolyData) - mesh = pv.StructuredGrid(x, y, z).extract_surface() + self.grid = pv.StructuredGrid(x, y, z) + mesh = self.grid.extract_surface() super().__init__(mesh=mesh, name=name, parent=parent, children=children, **kwargs) + def add_texture_cmap_along_axis( + self, axis: Optional[str] = None, cmap: Optional[Union[str, Colormap]] = None, n_colors: Optional[int] = None + ): + """Create mesh texture from a mathplotlib colormap and the variation along an axis. + + By default, the variation is along the Y axis (elevation). + + Parameters + ---------- + axis : str, optional + Axis along which to vary the colormap. + If None, the variation is along the Y axis (elevation). + + cmap : str or Colormap, optional + Colormap to use from matplotlib. + If None, the default colormap 'nipy_spectral' is used. -class ProcgenGrid(Object3D): + n_colors : int, optional + Number of colors to use in the colormap. + If None, the number of colors is the total number in the colormap. + """ + if cmap is None: + cmap = "nipy_spectral" + cmap_fct = get_cmap(name=cmap, lut=n_colors) + + if axis is None: + axis = "y" + + if axis == "x": + points = self.grid.x + elif axis == "y": + points = self.grid.y + elif axis == "z": + points = self.grid.z + else: + raise ValueError("axis must be one of x, y, z") + points = points.squeeze(-1) + x = points.ravel() + hue = (x - np.nanmin(x)) / (np.nanmax(x) - np.nanmin(x)) + colors = (cmap_fct(hue)[:, 0:3] * 255.0).astype(np.uint8) + image = colors.reshape((points.shape[0], points.shape[1], 3)) # [:-1, :-1, :] # , order="F") + + self.material.base_color_texture = pv.Texture(image) + + # Define the texture coordinates from the bounds of the grid + b = self.mesh.GetBounds() # [xmin, xmax, ymin, ymax, zmin, zmax] + origin = [b[0], b[2], b[4]] # [xmin, ymin, zmin] + point_u = [b[1], b[2], b[4]] # [xmax, ymin, zmin] + point_v = [b[0], b[2], b[5]] # [xmin, ymin, zmax] + self.mesh.texture_map_to_plane( + origin=origin, point_u=point_u, point_v=point_v, inplace=True + ) # Map the structure to the plane + + +class ProcgenGrid(StructuredGrid): """Create a procedural generated 3D grid (structured plane) from tiles / previous map. @@ -1180,6 +1262,9 @@ def __init__( else: self.map_2d = specific_map + raise NotImplementedError( + "Shallow generation not implemented yet or maybe not needed." + ) # TODO remove shallow gen? else: # Saves these for other functions that might use them # We take index 0 since generate_map is now vectorized, but we don't have @@ -1187,9 +1272,8 @@ def __init__( coordinates, map_2ds = generate_map(specific_map=specific_map, **all_args) self.coordinates, self.map_2d = coordinates[0], map_2ds[0] - # If it is a structured grid, extract the surface mesh (PolyData) - mesh = pv.StructuredGrid(*self.coordinates).extract_surface() - super().__init__(mesh=mesh, name=name, parent=parent, children=children, **kwargs) + x, y, z = self.coordinates + super().__init__(x=x, y=y, z=z, name=name, parent=parent, children=children, **kwargs) def generate_3D( self, @@ -1210,6 +1294,21 @@ def generate_3D( class ProcGenPrimsMaze3D(Asset): + """ + Procedurally generated 3D maze. + + Parameters + ---------- + width: int + Width of the maze. + + height: int + Height of the maze. + + depth: int + Depth of the maze. + """ + __NEW_ID = itertools.count() # Singleton to count instances of the classes for automatic naming def __init__(self, width: int, depth: int, name=None, wall_keep_prob=0.5, wall_material=None, **kwargs): diff --git a/src/simenv/assets/utils.py b/src/simenv/assets/utils.py index 3d37b7b6..a132866a 100644 --- a/src/simenv/assets/utils.py +++ b/src/simenv/assets/utils.py @@ -42,6 +42,16 @@ def snakecase_to_camelcase(name: str) -> str: return "".join(n.capitalize() for n in itertools.chain.from_iterable(name) if n != "") +class classproperty(object): + # required to use a classmethod as a property + # see https://stackoverflow.com/questions/128573/using-property-on-classmethods + def __init__(self, fget): + self.fget = fget + + def __get__(self, owner_self, owner_cls): + return self.fget(owner_cls) + + def get_transform_from_trs( translation: Union[np.ndarray, List[float]], rotation: Union[np.ndarray, List[float]], diff --git a/src/simenv/engine/pyvista_engine.py b/src/simenv/engine/pyvista_engine.py index b84d2ea8..145c7de9 100644 --- a/src/simenv/engine/pyvista_engine.py +++ b/src/simenv/engine/pyvista_engine.py @@ -68,9 +68,9 @@ def _view_vector(*args: Any) -> None: class PyVistaEngine(Engine): - def __init__(self, scene, auto_update=True, **plotter_kwargs): + def __init__(self, scene, auto_update=True, **add_mesh_kwargs): self.plotter: pyvista.Plotter = None - self.plotter_kwargs = plotter_kwargs + self.add_mesh_kwargs = add_mesh_kwargs self.auto_update = bool(CustomBackgroundPlotter is not None and auto_update) self._scene: Asset = scene @@ -78,7 +78,7 @@ def __init__(self, scene, auto_update=True, **plotter_kwargs): def _initialize_plotter(self): plotter_args = {"lighting": "none"} - plotter_args.update(self.plotter_kwargs) + plotter_args.update(self.add_mesh_kwargs) if self.auto_update: self.plotter: pyvista.Plotter = CustomBackgroundPlotter(**plotter_args) else: @@ -130,7 +130,7 @@ def update_asset(self, asset_node): self.plotter.reset_camera() - def _add_asset_to_scene(self, node, model_transform_matrix): + def _add_asset_to_scene(self, node, model_transform_matrix, **add_mesh_kwargs): if self.plotter is None or not hasattr(self.plotter, "ren_win"): return @@ -139,7 +139,7 @@ def _add_asset_to_scene(self, node, model_transform_matrix): located_mesh = node.mesh.transform(model_transform_matrix, inplace=False) # Material if node.material is None: - actor = self.plotter.add_mesh(located_mesh) + actor = self.plotter.add_mesh(located_mesh, **add_mesh_kwargs) else: material = node.material actor = self.plotter.add_mesh( @@ -152,8 +152,9 @@ def _add_asset_to_scene(self, node, model_transform_matrix): texture=None, # We set all the textures ourself in _set_pbr_material_for_actor specular_power=1.0, # Fixing a default of pyvista point_size=1.0, # Fixing a default of pyvista + **add_mesh_kwargs, ) - self._set_pbr_material_for_actor(actor, material) + self._set_pbr_material_for_actor(actor, material, located_mesh) self._plotter_actors[node.uuid] = actor @@ -169,7 +170,7 @@ def _add_asset_to_scene(self, node, model_transform_matrix): self._plotter_actors[node.uuid] = self.plotter.add_light(light) @staticmethod - def _set_pbr_material_for_actor(actor: pyvista._vtk.vtkActor, material: Material): + def _set_pbr_material_for_actor(actor: pyvista._vtk.vtkActor, material: Material, mesh: pyvista.DataSet): """Set all the necessary properties for a nice PBR material rendering Inspired by https://github.com/Kitware/VTK/blob/master/IO/Import/vtkGLTFImporter.cxx#L188 """ @@ -190,6 +191,12 @@ def _set_pbr_material_for_actor(actor: pyvista._vtk.vtkActor, material: Material actor.GetProperty().BackfaceCullingOn() if material.base_color_texture: + if mesh.GetPointData().GetTCoords() is None: + raise ValueError( + "This mesh doesn't have texture coordinates. " + "You need to define them to use a texture. " + "You can set texture coordinate with mesh.active_t_coords." + ) # set albedo texture material.base_color_texture.UseSRGBColorSpaceOn() prop.SetBaseColorTexture(material.base_color_texture) @@ -247,7 +254,7 @@ def _set_pbr_material_for_actor(actor: pyvista._vtk.vtkActor, material: Material actor.GetProperty().SetNormalScale(1.0) prop.SetNormalTexture(material.normal_texture) - def regenerate_scene(self): + def regenerate_scene(self, **add_mesh_kwargs): if self.plotter is None or not hasattr(self.plotter, "ren_win"): self._initialize_plotter() @@ -265,19 +272,25 @@ def regenerate_scene(self): else: model_transform_matrix = transforms[0] - self._add_asset_to_scene(node, model_transform_matrix) + self._add_asset_to_scene(node, model_transform_matrix, **add_mesh_kwargs) if not self.plotter.renderer.lights: self.plotter.enable_lightkit() # Still add some lights self.plotter.reset_camera() - def show(self, auto_update: Optional[bool] = None, **plotter_kwargs): + def show(self, auto_update: Optional[bool] = None, **kwargs): if auto_update is not None and auto_update != self.auto_update: self.plotter = None self.auto_update = auto_update - self.regenerate_scene() + auto_close = kwargs.pop("auto_close", True) + if isinstance(self.plotter, pyvista.Plotter): + plotter_kwargs = {"auto_close": auto_close} + else: + plotter_kwargs = {} + + self.regenerate_scene(**kwargs) self.plotter.show(**plotter_kwargs) def close(self):