Fix omniverse integration notebook

modules/omniverse/omniverse_integration.ipynb

@@ -484,7 +484,10 @@
    "source": [
     "## Convert 3D model contain all organs to USD format\n",
     "\n",
-    "[Universal Scene Description (OpenUSD)](https://openusd.org/release/index.html) is an extensible ecosystem of file formats, compositors, renderers, and other plugins for comprehensive 3D scene description."
+    "[Universal Scene Description (OpenUSD)](https://openusd.org/release/index.html) is an extensible ecosystem of file formats, compositors, renderers, and other plugins for comprehensive 3D scene description.\n",
+    "\n",
+    "Instead of using code to convert, you can also use the built-in converter in NVIDIA Omniverse to perform the conversion. For more details on using the CAD Converter extension, please refer to the official documentation:\n",
+    "https://docs.omniverse.nvidia.com/extensions/latest/ext_cad-converter/manual.html"
    ]
   },
   {
@@ -501,7 +504,7 @@
     }
    ],
    "source": [
-    "obj_filename = f\"{bundle_root}/datasets/monai/obj/all_organs.gltf\"\n",
+    "obj_filename = f\"{bundle_root}/datasets/monai/obj/all_organs_modified.gltf\"\n",
     "usd_filename = f\"{bundle_root}/datasets/monai/obj/all_organs.usd\"\n",
     "\n",
     "convert_mesh_to_usd(obj_filename, usd_filename)"
@@ -570,7 +573,11 @@
    "source": [
     "## Visualization in the Omniverse\n",
     "\n",
-    "Download the [NVIDIA Omniverse](https://www.nvidia.com/en-us/omniverse/) launcher to explore applications such as USD Composer for viewing and manipulating the OpenUSD output file.\n",
+    "Since the Omniverse Launcher has been deprecated, you can use USD Composer to view and manipulate the OpenUSD output files.\n",
+    "\n",
+    "To download the latest version of USD Composer, you can follow the instructions here to build and launch:\n",
+    "\n",
+    "https://github.com/NVIDIA-Omniverse/kit-app-template/tree/main/templates/apps/usd_composer#getting-started\n",
     "\n",
     "![omniverse](./omniverse.png)"
    ]

modules/omniverse/utility.py

@@ -14,10 +14,12 @@
 import os
 import vtk
 import json
-from pxr import Usd, UsdGeom, Gf, Sdf
+from pxr import Usd, UsdGeom, Gf, Sdf, UsdShade
 import trimesh
 import numpy as np
 import matplotlib.pyplot as plt
+import random
+import colorsys
 
 
 def convert_to_mesh(
@@ -186,30 +188,71 @@ def convert_mesh_to_usd(input_file, output_file):
 
     # Create a new USD stage
     stage = Usd.Stage.CreateNew(output_file)
+    root = UsdGeom.Xform.Define(stage, "/World")
+    stage.SetDefaultPrim(root.GetPrim())
+    materials_path = "/World/Materials"
+    UsdGeom.Scope.Define(stage, materials_path)
 
     # If the mesh is a Scene, process each geometry
     if isinstance(mesh, trimesh.Scene):
-        for name, geometry in mesh.geometry.items():
-            # Create a unique path for each mesh
-            mesh_path = f"/{name}"
-            usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
+        for node_name in mesh.graph.nodes:
+            if node_name == "world":
+                continue
+            geom_name = mesh.graph.get(node_name)[1]
+            if geom_name is not None and geom_name.startswith("mesh"):
+                print(f"Processing mesh: {node_name} {geom_name}")
+                # Create a unique path for each mesh
+                node_path = f"/World/{node_name}"
+                xform = UsdGeom.Xform.Define(stage, node_path)
+                # Define the Mesh under the Xform
+                mesh_path = f"{node_path}/Mesh"
+                usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
+                # get the geometry of the node
+                geometry = mesh.geometry[geom_name]
+
+                # Create a random color for this mesh
+                # Using HSV for better color distribution
+                h = random.random()  # Random hue
+                s = 0.7 + 0.3 * random.random()  # Saturation between 0.7-1.0
+                v = 0.7 + 0.3 * random.random()  # Value between 0.7-1.0
+                r, g, b = colorsys.hsv_to_rgb(h, s, v)
+
+                # Create a material with the random color
+                mat_name = f"{node_name}_material"
+                mat_path = f"{materials_path}/{mat_name}"
+                material = UsdShade.Material.Define(stage, mat_path)
+
+                # Create shader
+                shader = UsdShade.Shader.Define(stage, f"{mat_path}/PreviewSurface")
+                shader.CreateIdAttr("UsdPreviewSurface")
+
+                # Set the random color
+                shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(r, g, b))
+                shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.4)
+
+                # Connect shader to material
+                material_output = material.CreateOutput("surface", Sdf.ValueTypeNames.Token)
+                shader_output = shader.CreateOutput("surface", Sdf.ValueTypeNames.Token)
+                material_output.ConnectToSource(shader_output)
+
+                # Bind material to mesh
+                UsdShade.MaterialBindingAPI(usd_mesh).Bind(material)
+
+                # Set vertex positions
+                usd_mesh.GetPointsAttr().Set([Gf.Vec3f(*vertex) for vertex in geometry.vertices])
+
+                # Set face indices and counts
+                face_vertex_indices = geometry.faces.flatten().tolist()
+                face_vertex_counts = [len(face) for face in geometry.faces]
+
+                usd_mesh.GetFaceVertexIndicesAttr().Set(face_vertex_indices)
+                usd_mesh.GetFaceVertexCountsAttr().Set(face_vertex_counts)
+
+                # Optionally, set normals
+                if geometry.vertex_normals is not None:
+                    usd_mesh.GetNormalsAttr().Set([Gf.Vec3f(*normal) for normal in geometry.vertex_normals])
+                    usd_mesh.SetNormalsInterpolation("vertex")
 
-            # Set vertex positions
-            usd_mesh.GetPointsAttr().Set([Gf.Vec3f(*vertex) for vertex in geometry.vertices])
-
-            # Set face indices and counts
-            face_vertex_indices = geometry.faces.flatten().tolist()
-            face_vertex_counts = [len(face) for face in geometry.faces]
-
-            usd_mesh.GetFaceVertexIndicesAttr().Set(face_vertex_indices)
-            usd_mesh.GetFaceVertexCountsAttr().Set(face_vertex_counts)
-
-            # Optionally, set normals
-            if geometry.vertex_normals is not None:
-                usd_mesh.GetNormalsAttr().Set([Gf.Vec3f(*normal) for normal in geometry.vertex_normals])
-                usd_mesh.SetNormalsInterpolation("vertex")
-
-            # Handle materials and other attributes if needed
     else:
         # It's a single mesh, proceed as before
         usd_mesh = UsdGeom.Mesh.Define(stage, "/Mesh")