From 8223b2233631178c499660593d89704770754817 Mon Sep 17 00:00:00 2001 From: Thomwolf Date: Fri, 29 Jul 2022 14:30:49 +0200 Subject: [PATCH 1/7] adding 1d texture example for grid --- .gitignore | 3 ++ examples/procgen_grid.py | 15 +++++- src/simenv/assets/__init__.py | 2 + src/simenv/assets/colors.py | 76 +++++++++++++++++++++++++++++ src/simenv/assets/material.py | 43 +++++++--------- src/simenv/assets/object.py | 47 ++++++++++++++++-- src/simenv/assets/textures.py | 35 +++++++++++++ src/simenv/assets/utils.py | 10 ++++ src/simenv/engine/pyvista_engine.py | 10 +++- 9 files changed, 208 insertions(+), 33 deletions(-) create mode 100644 src/simenv/assets/colors.py create mode 100644 src/simenv/assets/textures.py diff --git a/.gitignore b/.gitignore index c68546c9..fae89ecd 100644 --- a/.gitignore +++ b/.gitignore @@ -146,3 +146,6 @@ wfc_binding.cpp .gen_files/ .config/ + +# Mac +*.icloud \ No newline at end of file diff --git a/examples/procgen_grid.py b/examples/procgen_grid.py index f2cf0a97..707bf235 100644 --- a/examples/procgen_grid.py +++ b/examples/procgen_grid.py @@ -19,8 +19,19 @@ ) * 0.6 ) -scene += sm.ProcgenGrid(specific_map=specific_map) -scene += sm.Light() +proc_grid = sm.ProcgenGrid(specific_map=specific_map) + +# Let's color our grid by height. +# We assign a simple 1D texture to the object and assing mesh texture coordinates from the height of each points of the mesh. +proc_grid.material.base_color_texture = sm.TEXTURE_1D_CMAP # Set a simple 1D color map texture +height = proc_grid.mesh.points[:, 1, None] # We get the height coordinates of the mesh points ( theY axis) +texture_coord = np.concatenate( + [height, np.zeros_like(height)], axis=1 +) # Make texture coordinates (UV) by adding a zeros axis +proc_grid.mesh.active_t_coords = texture_coord # Assign as point to texture mapping + +scene += proc_grid +scene += sm.LightSun() scene.show() input("Press Enter for second scene") diff --git a/src/simenv/assets/__init__.py b/src/simenv/assets/__init__.py index 3f9ca2fa..5218cb52 100644 --- a/src/simenv/assets/__init__.py +++ b/src/simenv/assets/__init__.py @@ -1,8 +1,10 @@ from .asset import * from .camera import * from .collider import * +from .colors import * from .light import * from .material import * from .object import * from .rigidbody_component import * +from .textures import * from .utils import * diff --git a/src/simenv/assets/colors.py b/src/simenv/assets/colors.py new file mode 100644 index 00000000..7642ed18 --- /dev/null +++ b/src/simenv/assets/colors.py @@ -0,0 +1,76 @@ +# 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())) +CMAP_ONLY_COLORS = np.array(list(COLORS_NO_GRAYSCALE.values())) +CMAP_ONLY_GRAYSCALE = np.array(list(COLORS_ONLY_GRAYSCALE.values())) diff --git a/src/simenv/assets/material.py b/src/simenv/assets/material.py index 56e3d20b..6ea7a4ef 100644 --- a/src/simenv/assets/material.py +++ b/src/simenv/assets/material.py @@ -22,17 +22,8 @@ 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 @@ -167,60 +158,60 @@ 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) diff --git a/src/simenv/assets/object.py b/src/simenv/assets/object.py index 02c0f1a2..fd87d1fa 100644 --- a/src/simenv/assets/object.py +++ b/src/simenv/assets/object.py @@ -46,7 +46,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, @@ -57,14 +57,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. @@ -1207,6 +1233,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/textures.py b/src/simenv/assets/textures.py new file mode 100644 index 00000000..b13d0c39 --- /dev/null +++ b/src/simenv/assets/textures.py @@ -0,0 +1,35 @@ +# 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 shortcut for textures.""" +import numpy as np +import pyvista as pv + +from .colors import CMAP_ONLY_COLORS + + +# class Texture(pv.Texture): +# def __init__(self, *args, **kwargs): +# super().__init__(*args, **kwargs) + +# @classproperty +# def TEXTURE_1D_CMAP(cls) -> "Texture": +# return cls.from_pyvista(pv.numpy_to_texture(CMAP_ONLY_COLORS[:, np.newaxis, :])) + +TEXTURE_1D_CMAP = pv.numpy_to_texture(CMAP_ONLY_COLORS[:, np.newaxis, :]) + + +def texture_from_1d_cmap(np_array: np.array) -> pv.Texture: + return pv.numpy_to_texture(np_array[:, np.newaxis, :]) 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 dc27a88c..8503d9a3 100644 --- a/src/simenv/engine/pyvista_engine.py +++ b/src/simenv/engine/pyvista_engine.py @@ -155,7 +155,7 @@ def _add_asset_to_scene(self, node, model_transform_matrix): specular_power=1.0, # Fixing a default of pyvista point_size=1.0, # Fixing a default of pyvista ) - self._set_pbr_material_for_actor(actor, material) + self._set_pbr_material_for_actor(actor, material, located_mesh) self._plotter_actors[node.uuid] = actor @@ -171,7 +171,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 """ @@ -192,6 +192,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) From b97dbd609c91cf0015e16630c70f13d61c320001 Mon Sep 17 00:00:00 2001 From: Thomwolf Date: Fri, 29 Jul 2022 15:23:29 +0200 Subject: [PATCH 2/7] fix style --- src/simenv/assets/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/simenv/assets/__init__.py b/src/simenv/assets/__init__.py index 2acf8419..21682029 100644 --- a/src/simenv/assets/__init__.py +++ b/src/simenv/assets/__init__.py @@ -6,6 +6,6 @@ from .material import * from .object import * from .rigidbody_component import * -from .textures import * from .sensors import CameraSensor, RaycastSensor, Sensor, StateSensor +from .textures import * from .utils import * From f40d540e26c6211a03763075ea4c62d62b01ae30 Mon Sep 17 00:00:00 2001 From: Nathan Lambert Date: Wed, 3 Aug 2022 15:37:54 -0700 Subject: [PATCH 3/7] add greyscale --- src/simenv/assets/textures.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/simenv/assets/textures.py b/src/simenv/assets/textures.py index b13d0c39..f30467e7 100644 --- a/src/simenv/assets/textures.py +++ b/src/simenv/assets/textures.py @@ -17,7 +17,7 @@ import numpy as np import pyvista as pv -from .colors import CMAP_ONLY_COLORS +from .colors import CMAP_ONLY_COLORS, CMAP_ONLY_GRAYSCALE # class Texture(pv.Texture): @@ -29,6 +29,7 @@ # return cls.from_pyvista(pv.numpy_to_texture(CMAP_ONLY_COLORS[:, np.newaxis, :])) TEXTURE_1D_CMAP = pv.numpy_to_texture(CMAP_ONLY_COLORS[:, np.newaxis, :]) +TEXTURE_1D_CMAP_GRAY = pv.numpy_to_texture(CMAP_ONLY_GRAYSCALE[:, np.newaxis, :]) def texture_from_1d_cmap(np_array: np.array) -> pv.Texture: From ba5982906a5a87d0ce9d9cf0dbdeb89061745d01 Mon Sep 17 00:00:00 2001 From: Thomwolf Date: Sun, 7 Aug 2022 00:55:03 +0200 Subject: [PATCH 4/7] fixing textures for structuredgrid --- examples/procgen_grid.py | 24 +++++++------ src/simenv/assets/__init__.py | 1 - src/simenv/assets/colors.py | 18 ++++++++-- src/simenv/assets/material.py | 31 +++++++++++++---- src/simenv/assets/object.py | 54 ++++++++++++++++++++++++++--- src/simenv/assets/textures.py | 36 ------------------- src/simenv/engine/pyvista_engine.py | 22 ++++++------ 7 files changed, 115 insertions(+), 71 deletions(-) delete mode 100644 src/simenv/assets/textures.py diff --git a/examples/procgen_grid.py b/examples/procgen_grid.py index 9f060725..cc6fd3b9 100644 --- a/examples/procgen_grid.py +++ b/examples/procgen_grid.py @@ -21,19 +21,13 @@ proc_grid = sm.ProcgenGrid(specific_map=specific_map) -# Let's color our grid by height. -# We assign a simple 1D texture to the object and assing mesh texture coordinates from the height of each points of the mesh. -proc_grid.material.base_color_texture = sm.TEXTURE_1D_CMAP # Set a simple 1D color map texture -height = proc_grid.mesh.points[:, 1, None] # We get the height coordinates of the mesh points ( theY axis) -texture_coord = np.concatenate( - [height, np.zeros_like(height)], axis=1 -) # Make texture coordinates (UV) by adding a zeros axis -proc_grid.mesh.active_t_coords = texture_coord # Assign as point to texture mapping - +# 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") scene.close() @@ -41,6 +35,15 @@ # 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() @@ -60,6 +63,7 @@ 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") diff --git a/src/simenv/assets/__init__.py b/src/simenv/assets/__init__.py index 21682029..902d6694 100644 --- a/src/simenv/assets/__init__.py +++ b/src/simenv/assets/__init__.py @@ -7,5 +7,4 @@ from .object import * from .rigidbody_component import * from .sensors import CameraSensor, RaycastSensor, Sensor, StateSensor -from .textures import * from .utils import * diff --git a/src/simenv/assets/colors.py b/src/simenv/assets/colors.py index 7642ed18..3c9c7a1a 100644 --- a/src/simenv/assets/colors.py +++ b/src/simenv/assets/colors.py @@ -71,6 +71,18 @@ "GRAY100": BLACK, } -CMAP_ALL = np.array(list(COLORS_ALL.values())) -CMAP_ONLY_COLORS = np.array(list(COLORS_NO_GRAYSCALE.values())) -CMAP_ONLY_GRAYSCALE = np.array(list(COLORS_ONLY_GRAYSCALE.values())) +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 3ce08f9f..62c7b94a 100644 --- a/src/simenv/assets/material.py +++ b/src/simenv/assets/material.py @@ -17,7 +17,7 @@ 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 @@ -92,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 @@ -109,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): @@ -216,3 +231,7 @@ def PURPLE(cls): @classproperty def OLIVE(cls): 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 fd87d1fa..e9587e06 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 @@ -1080,11 +1082,51 @@ 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. -class ProcgenGrid(Object3D): + By default, the variation is along the Y axis (elevation). + """ + 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. @@ -1203,6 +1245,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 @@ -1210,9 +1255,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, diff --git a/src/simenv/assets/textures.py b/src/simenv/assets/textures.py deleted file mode 100644 index f30467e7..00000000 --- a/src/simenv/assets/textures.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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 shortcut for textures.""" -import numpy as np -import pyvista as pv - -from .colors import CMAP_ONLY_COLORS, CMAP_ONLY_GRAYSCALE - - -# class Texture(pv.Texture): -# def __init__(self, *args, **kwargs): -# super().__init__(*args, **kwargs) - -# @classproperty -# def TEXTURE_1D_CMAP(cls) -> "Texture": -# return cls.from_pyvista(pv.numpy_to_texture(CMAP_ONLY_COLORS[:, np.newaxis, :])) - -TEXTURE_1D_CMAP = pv.numpy_to_texture(CMAP_ONLY_COLORS[:, np.newaxis, :]) -TEXTURE_1D_CMAP_GRAY = pv.numpy_to_texture(CMAP_ONLY_GRAYSCALE[:, np.newaxis, :]) - - -def texture_from_1d_cmap(np_array: np.array) -> pv.Texture: - return pv.numpy_to_texture(np_array[:, np.newaxis, :]) diff --git a/src/simenv/engine/pyvista_engine.py b/src/simenv/engine/pyvista_engine.py index 8503d9a3..d5553ba7 100644 --- a/src/simenv/engine/pyvista_engine.py +++ b/src/simenv/engine/pyvista_engine.py @@ -70,9 +70,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 @@ -80,7 +80,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: @@ -132,7 +132,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 @@ -141,7 +141,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( @@ -154,6 +154,7 @@ 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, located_mesh) @@ -255,7 +256,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() @@ -273,20 +274,21 @@ 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() - self.plotter.show(**plotter_kwargs) + auto_close = kwargs.pop("auto_close", True) + self.regenerate_scene(**kwargs) + self.plotter.show(auto_close=auto_close) def close(self): if self.plotter is not None: From 515e346a630f611d4eefa9155f6127e03402a828 Mon Sep 17 00:00:00 2001 From: Thomwolf Date: Sun, 7 Aug 2022 09:14:19 +0200 Subject: [PATCH 5/7] more docstring --- src/simenv/assets/object.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/simenv/assets/object.py b/src/simenv/assets/object.py index 7c1ec1b2..0ca788e0 100644 --- a/src/simenv/assets/object.py +++ b/src/simenv/assets/object.py @@ -1091,6 +1091,20 @@ def add_texture_cmap_along_axis( """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. + + 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" From 2e433836b0e51db055c37b1a339c6abb647f80da Mon Sep 17 00:00:00 2001 From: Thomwolf Date: Mon, 8 Aug 2022 15:51:28 +0200 Subject: [PATCH 6/7] fix --- examples/procgen_grid.py | 6 +++--- src/simenv/assets/object.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/procgen_grid.py b/examples/procgen_grid.py index 572978b4..355ea3c5 100644 --- a/examples/procgen_grid.py +++ b/examples/procgen_grid.py @@ -29,7 +29,7 @@ scene += sm.LightSun() scene.show(show_edges=True) -input("Press Enter for second scene") +# input("Press Enter for second scene") scene.close() scene.clear() @@ -47,7 +47,7 @@ scene += sm.LightSun() scene.show() -input("Press Enter for third scene") +# input("Press Enter for third scene") scene.close() scene.clear() @@ -66,6 +66,6 @@ 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/object.py b/src/simenv/assets/object.py index 0ca788e0..59b58c1d 100644 --- a/src/simenv/assets/object.py +++ b/src/simenv/assets/object.py @@ -1125,7 +1125,7 @@ def add_texture_cmap_along_axis( 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") + image = colors.reshape((points.shape[0], points.shape[1], 3)) # [:-1, :-1, :] # , order="F") self.material.base_color_texture = pv.Texture(image) From fa3de016ec7759b2f74494a08d8139fd2bf62325 Mon Sep 17 00:00:00 2001 From: Thomwolf Date: Mon, 8 Aug 2022 16:00:58 +0200 Subject: [PATCH 7/7] fix auto-close --- src/simenv/engine/pyvista_engine.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/simenv/engine/pyvista_engine.py b/src/simenv/engine/pyvista_engine.py index 5036c198..145c7de9 100644 --- a/src/simenv/engine/pyvista_engine.py +++ b/src/simenv/engine/pyvista_engine.py @@ -285,8 +285,13 @@ def show(self, auto_update: Optional[bool] = None, **kwargs): self.auto_update = auto_update 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(auto_close=auto_close) + self.plotter.show(**plotter_kwargs) def close(self): if self.plotter is not None: