diff --git a/aitviewer/renderables/lines.py b/aitviewer/renderables/lines.py index 1d32456..be14441 100644 --- a/aitviewer/renderables/lines.py +++ b/aitviewer/renderables/lines.py @@ -18,6 +18,7 @@ compute_vertex_and_face_normals, set_lights_in_program, set_material_properties, + usd, ) from aitviewer.utils.decorators import hooked from aitviewer.utils.so3 import aa2rot_numpy as aa2rot @@ -432,6 +433,47 @@ def remove_frames(self, frames): self.lines = np.delete(self.lines, frames, axis=0) self.redraw() + def export_usd(self, stage, usd_path: str, directory: str = None, verbose=False): + name = f"{self.name}_{self.uid:03}".replace(" ", "_") + usd_path = f"{usd_path}/{name}" + + if self.mode == "lines": + v0s = self.lines[:, ::2] + v1s = self.lines[:, 1::2] + else: + v0s = self.lines[:, :-1] + v1s = self.lines[:, 1:] + + print(self.lines.shape) + print(v0s.shape) + + # Data is in the form of (F, N_LINES, 3), convert it to (F*N_LINES, 3) + v0s = np.reshape(v0s, (-1, 3)) + v1s = np.reshape(v1s, (-1, 3)) + + self.r_tip = self.r_base if self.r_tip is None else self.r_tip + + # If r_tip is below a certain threshold, we create a proper cone, i.e. with just a single vertex at the top. + if self.r_tip < 10e-6: + data = _create_cone_from_to(v0s, v1s, radius=self.r_base) + else: + data = _create_cylinder_from_to(v0s, v1s, radius1=self.r_base, radius2=self.r_tip) + + L = self.n_lines + V = data["vertices"].shape[1] + + vertices = data["vertices"].reshape((self.n_frames, -1, 3)) + faces = data["faces"] + + fs = faces[np.newaxis].repeat(L, 0).reshape((L, -1)) + offsets = (np.arange(L) * V).reshape((L, 1)) + faces = (fs + offsets).reshape((-1, 3)) + + mesh = usd.add_mesh(stage, usd_path, self.name, vertices, faces, self.get_local_transform()) + usd.add_color(stage, mesh, usd_path, self.color[:3]) + + self._export_usd_recursively(stage, usd_path, directory, verbose) + class Lines2D(Node): """Render 2D lines.""" diff --git a/aitviewer/renderables/meshes.py b/aitviewer/renderables/meshes.py index 8dba0f3..4e3161c 100644 --- a/aitviewer/renderables/meshes.py +++ b/aitviewer/renderables/meshes.py @@ -746,19 +746,7 @@ def export_usd(self, stage, usd_path: str, directory: str = None, verbose=False) name = f"{self.name}_{self.uid:03}".replace(" ", "_") usd_path = f"{usd_path}/{name}" - # Transform. - xform = UsdGeom.Xform.Define(stage, usd_path) - a_xform = xform.AddTransformOp() - a_xform.Set(Gf.Matrix4d(self.get_local_transform().astype(np.float64).T)) - - # Geometry. - mesh = UsdGeom.Mesh.Define(stage, usd_path + "/" + self.name.replace(" ", "_")) - a_vertices = mesh.CreatePointsAttr() - for i in range(self.n_frames): - a_vertices.Set(time=i + 1, value=self.vertices[i]) - mesh.CreateFaceVertexCountsAttr(np.full(self.faces.shape[0], 3)) - mesh.CreateFaceVertexIndicesAttr(self.faces) - + mesh = usd.add_mesh(stage, usd_path, self.name, self.vertices, self.faces, self.get_local_transform()) if self.has_texture and not self.use_pickle_texture: # UVs. a_uv = UsdGeom.PrimvarsAPI(mesh).CreatePrimvar( @@ -771,6 +759,30 @@ def export_usd(self, stage, usd_path: str, directory: str = None, verbose=False) else: texture_path = usd.copy_texture(self.texture_path, name, directory) usd.add_texture(stage, mesh, usd_path, texture_path) + else: + # NOTE: Per vertex and per face colors using usd displayColor are not currently + # loaded by Blender. This code path can be enabled once support is there. + if False: + a_colors = mesh.GetDisplayColorAttr() + if self._face_colors is not None: + # Per face colors. + if self._face_colors.shape[0] == 1: + a_colors.Set(self._face_colors[0, :, :3].astype(np.float32)) + else: + for i in range(self.n_frames): + a_colors.Set(time=i + 1, value=self._face_colors[i, :, :3].astype(np.float32)) + elif self._vertex_colors is not None: + # Per vertex colors. + if self._vertex_colors.shape[0] == 1: + a_colors.Set(self._vertex_colors[0, :, :3].astype(np.float32)) + else: + for i in range(self.n_frames): + a_colors.Set(time=i + 1, value=self._vertex_colors[i, :, :3].astype(np.float32)) + else: + # Uniform color. + a_colors.Set(np.array(self.color, np.float32)[:3]) + else: + usd.add_color(stage, mesh, usd_path, self.color[:3]) self._export_usd_recursively(stage, usd_path, directory, verbose) diff --git a/aitviewer/renderables/spheres.py b/aitviewer/renderables/spheres.py index d3e211b..5c09392 100644 --- a/aitviewer/renderables/spheres.py +++ b/aitviewer/renderables/spheres.py @@ -11,6 +11,7 @@ get_outline_program, get_sphere_instanced_program, ) +from aitviewer.utils import usd from aitviewer.utils.decorators import hooked from aitviewer.utils.utils import set_lights_in_program, set_material_properties @@ -40,7 +41,7 @@ def _create_sphere(radius=1.0, rings=16, sectors=32): v += 1 n += 1 - faces = np.zeros([rings * sectors * 2, 3], dtype=np.int32) + faces = np.zeros([(rings - 1) * (sectors - 1) * 2, 3], dtype=np.int32) i = 0 for r in range(rings - 1): for s in range(sectors - 1): @@ -264,6 +265,28 @@ def remove_frames(self, frames): self.sphere_positions = np.delete(self.sphere_positions, frames, axis=0) self.redraw() + def export_usd(self, stage, usd_path: str, directory: str = None, verbose=False): + name = f"{self.name}_{self.uid:03}".replace(" ", "_") + usd_path = f"{usd_path}/{name}" + + V = self.vertices.shape[0] + N = self.sphere_positions.shape[0] + M = self.n_spheres + + vertices = np.empty((N, V * M, 3), np.float32) + for i in range(N): + vs = self.vertices[np.newaxis].repeat(M, 0) + vertices[i] = (vs * self.radius + self.sphere_positions[i].reshape(M, 1, 3)).reshape((-1, 3)) + + fs = self.faces[np.newaxis].repeat(M, 0).reshape((M, -1)) + offsets = (np.arange(M) * V).reshape((M, 1)) + faces = (fs + offsets).reshape((-1, 3)) + + mesh = usd.add_mesh(stage, usd_path, self.name, vertices, faces, self.get_local_transform()) + usd.add_color(stage, mesh, usd_path, self.color[:3]) + + self._export_usd_recursively(stage, usd_path, directory, verbose) + class SpheresTrail(Spheres): """A sequence of spheres that leaves a trail, i.e. the past spheres keep being rendered.""" diff --git a/aitviewer/scene/light.py b/aitviewer/scene/light.py index 9bc01f8..d030a9f 100644 --- a/aitviewer/scene/light.py +++ b/aitviewer/scene/light.py @@ -2,11 +2,11 @@ from functools import lru_cache import numpy as np +from pxr import Gf, UsdGeom, UsdLux from aitviewer.renderables.lines import Lines from aitviewer.renderables.rigid_bodies import RigidBodies from aitviewer.scene.camera_utils import look_at, orthographic_projection -from aitviewer.scene.material import Material from aitviewer.scene.node import Node from aitviewer.utils.utils import ( direction_from_spherical_coordinates, @@ -58,6 +58,7 @@ def __init__( self.mesh.spheres.material.diffuse = 0.0 self.mesh.spheres.material.ambient = 1.0 self.mesh.spheres.color = (*tuple(light_color), 1.0) + self.mesh.export_usd_enabled = False self.add(self.mesh, show_in_hierarchy=False, enabled=False) @classmethod @@ -170,6 +171,7 @@ def _update_debug_lines(self): if self._debug_lines is None: self._debug_lines = Lines(lines, r_base=0.05, mode="lines", cast_shadow=False, is_selectable=False) + self._debug_lines.export_usd_enabled = False self.add(self._debug_lines, show_in_hierarchy=False) else: self._debug_lines.lines = lines @@ -281,3 +283,20 @@ def gui(self, imgui): else: if self._debug_lines: self._debug_lines.enabled = False + + def export_usd(self, stage, usd_path: str, directory: str = None, verbose=False): + name = f"{self.name}_{self.uid:03}".replace(" ", "_") + usd_path = f"{usd_path}/{name}" + + # Transform. + xform = UsdGeom.Xform.Define(stage, usd_path) + a_xform = xform.AddTransformOp() + a_xform.Set(Gf.Matrix4d(self.get_local_transform().astype(np.float64).T)) + + # Light. + light = UsdLux.DistantLight.Define(stage, usd_path + "/" + name.replace(" ", "_")) + lightAPI = light.LightAPI() + lightAPI.GetInput("color").Set(self.light_color[:3]) + lightAPI.GetInput("intensity").Set(self.strength) + + self._export_usd_recursively(stage, usd_path, directory, verbose) diff --git a/aitviewer/scene/node.py b/aitviewer/scene/node.py index 6b09ea6..101b251 100644 --- a/aitviewer/scene/node.py +++ b/aitviewer/scene/node.py @@ -712,7 +712,7 @@ def _export_usd_recursively(self, stage, usd_path, directory, verbose): if verbose: print(usd_path) for n in self.nodes: - if n.export_usd_enabled and n.show_in_hierarchy: + if n.export_usd_enabled: n.export_usd(stage, usd_path, directory, verbose) def export_usd(self, stage, usd_path: str, directory: str = None, verbose=False): diff --git a/aitviewer/scene/scene.py b/aitviewer/scene/scene.py index 8ccb3f7..855c09f 100644 --- a/aitviewer/scene/scene.py +++ b/aitviewer/scene/scene.py @@ -82,6 +82,7 @@ def __init__(self, **kwargs): color=(0.2, 0.2, 0.2, 1), mode="lines", ) + self.camera_target.export_usd_enabled = False self.add(self.camera_target, show_in_hierarchy=False, enabled=False) # Camera trackball. @@ -112,6 +113,7 @@ def __init__(self, **kwargs): trackball_colors, mode="lines", ) + self.trackball.export_usd_enabled = False self.add(self.trackball, show_in_hierarchy=False, enabled=False) self.custom_font = None diff --git a/aitviewer/utils/usd.py b/aitviewer/utils/usd.py index e26d333..61575ca 100644 --- a/aitviewer/utils/usd.py +++ b/aitviewer/utils/usd.py @@ -4,7 +4,7 @@ import numpy as np from PIL import Image -from pxr import Sdf, UsdShade +from pxr import Gf, Sdf, UsdGeom, UsdShade def _get_texture_paths(path, name, directory): @@ -28,6 +28,26 @@ def save_image_as_texture(img, img_name, name, directory): return usd_path +def add_color(stage, mesh, usd_path, color): + # Material. + mat_path = usd_path + "/material" + material = UsdShade.Material.Define(stage, mat_path) + + # Shader. + shader = UsdShade.Shader.Define(stage, mat_path + "/shader") + shader.CreateIdAttr("UsdPreviewSurface") + + # Connect the material to the shader. + material.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface") + + # Create a uniform color.np. + shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(tuple(color)) + + # Bind the Material to the mesh. + mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI) + UsdShade.MaterialBindingAPI(mesh).Bind(material) + + def add_texture(stage, mesh, usd_path, texture_path): # Material. mat_path = usd_path + "/material" @@ -58,3 +78,20 @@ def add_texture(stage, mesh, usd_path, texture_path): # Bind the Material to the mesh. mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI) UsdShade.MaterialBindingAPI(mesh).Bind(material) + + +def add_mesh(stage, usd_path, name, vertices, faces, transform): + # Transform. + xform = UsdGeom.Xform.Define(stage, usd_path) + a_xform = xform.AddTransformOp() + a_xform.Set(Gf.Matrix4d(transform.astype(np.float64).T)) + + # Geometry. + mesh = UsdGeom.Mesh.Define(stage, usd_path + "/" + name.replace(" ", "_")) + a_vertices = mesh.CreatePointsAttr() + for i in range(vertices.shape[0]): + a_vertices.Set(time=i + 1, value=vertices[i]) + mesh.CreateFaceVertexCountsAttr(np.full(faces.shape[0], 3)) + mesh.CreateFaceVertexIndicesAttr(faces) + + return mesh diff --git a/aitviewer/viewer.py b/aitviewer/viewer.py index 1a7038d..779beb8 100644 --- a/aitviewer/viewer.py +++ b/aitviewer/viewer.py @@ -1344,7 +1344,7 @@ def tree(nodes): imgui.end_child() if self.export_usd_name is None: - self.export_usd_name = "frame_{:0>6}".format(self.scene.current_frame_id) + self.export_usd_name = "scene" # HACK: we need to set the focus twice when the modal is first opened for it to take effect. if self._modal_focus_count > 0: @@ -1384,7 +1384,8 @@ def tree(nodes): self.export_usd( os.path.join(C.export_dir, "usd", f"{self.export_usd_name}"), self.export_usd_directory, - True, + False, + False, ) imgui.close_current_popup() self._export_usd_popup_open = False @@ -1772,20 +1773,21 @@ def export_frame(self, file_path, scale_factor: float = None, transparent_backgr self.run_animations = run_animations self._last_frame_rendered_at = self.timer.time - def export_usd(self, path: str, export_as_directory=False, verbose=False): + def export_usd(self, path: str, export_as_directory=False, verbose=False, ascii=False): from pxr import Usd, UsdGeom + extension = ".usd" if not ascii else ".usda" if export_as_directory: - if path.endswith(".usd"): + if path.endswith(extension): directory = path[:-4] else: directory = path os.makedirs(directory, exist_ok=True) - path = os.path.join(directory, os.path.basename(directory) + ".usd") + path = os.path.join(directory, os.path.basename(directory) + extension) else: directory = None - if not path.endswith(".usd"): - path += ".usd" + if not path.endswith(extension): + path += extension # Create a new file and setup scene parameters. stage = Usd.Stage.CreateNew(path)