USD export for geometric primitives and lights

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."""

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)
 

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."""

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)

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):

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

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

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)