Support both visual and collision meshes when loading URDF (nerfstudio-project#487)

* Display visual and collision meshes when loading URDF

* PR comments

* Bump min version of robot_descriptions

* Fix type annotation for trimesh scene

* Revert 09_urdf_visualizer.py

* Remove note about PR

* Try fix trimesh annotation

* pyright doesnt seem to know how to lint zip with strict kwarg

* Sort imports

* Handle some more edge cases

* nits

* incorporate feedback from @sea-bass

* example nit

---------

Co-authored-by: Brent Yi <yibrenth@gmail.com>

Author: sea-bass

examples_dev/09_urdf_visualizer.py

@@ -34,7 +34,7 @@ def create_robot_control_sliders(
     ) in viser_urdf.get_actuated_joint_limits().items():
         lower = lower if lower is not None else -np.pi
         upper = upper if upper is not None else np.pi
-        initial_pos = 0.0 if lower < 0 and upper > 0 else (lower + upper) / 2.0
+        initial_pos = 0.0 if lower < -0.1 and upper > 0.1 else (lower + upper) / 2.0
         slider = server.gui.add_slider(
             label=joint_name,
             min=lower,
@@ -66,16 +66,28 @@ def main(
         "anymal_c",
         "go2",
     ] = "panda",
+    load_meshes: bool = True,
+    load_collision_meshes: bool = False,
 ) -> None:
     # Start viser server.
     server = viser.ViserServer()
 
     # Load URDF.
     #
     # This takes either a yourdfpy.URDF object or a path to a .urdf file.
+    urdf = load_robot_description(
+        robot_type + "_description",
+        load_meshes=load_meshes,
+        build_scene_graph=load_meshes,
+        load_collision_meshes=load_collision_meshes,
+        build_collision_scene_graph=load_collision_meshes,
+    )
     viser_urdf = ViserUrdf(
         server,
-        urdf_or_path=load_robot_description(robot_type + "_description"),
+        urdf_or_path=urdf,
+        load_meshes=load_meshes,
+        load_collision_meshes=load_collision_meshes,
+        collision_mesh_color_override=(1.0, 0.0, 0.0, 0.5),
     )
 
     # Create sliders in GUI that help us move the robot joints.
@@ -84,10 +96,33 @@ def main(
             server, viser_urdf
         )
 
+    # Add visibility checkboxes.
+    with server.gui.add_folder("Visibility"):
+        show_meshes_cb = server.gui.add_checkbox(
+            "Show meshes",
+            viser_urdf.show_visual,
+        )
+        show_collision_meshes_cb = server.gui.add_checkbox(
+            "Show collision meshes", viser_urdf.show_collision
+        )
+
+    @show_meshes_cb.on_update
+    def _(_):
+        viser_urdf.show_visual = show_meshes_cb.value
+
+    @show_collision_meshes_cb.on_update
+    def _(_):
+        viser_urdf.show_collision = show_collision_meshes_cb.value
+
+    # Hide checkboxes if meshes are not loaded.
+    show_meshes_cb.visible = load_meshes
+    show_collision_meshes_cb.visible = load_collision_meshes
+
     # Set initial robot configuration.
     viser_urdf.update_cfg(np.array(initial_config))
 
     # Create grid.
+    trimesh_scene = viser_urdf._urdf.scene or viser_urdf._urdf.collision_scene
     server.scene.add_grid(
         "/grid",
         width=2,
@@ -96,7 +131,7 @@ def main(
             0.0,
             0.0,
             # Get the minimum z value of the trimesh scene.
-            viser_urdf._urdf.scene.bounds[0, 2],
+            trimesh_scene.bounds[0, 2] if trimesh_scene is not None else 0.0,
         ),
     )
 

pyproject.toml

@@ -62,7 +62,7 @@ examples = [
     "torch>=1.13.1",
     "matplotlib>=3.7.1",
     "plotly>=5.21.0",
-    "robot_descriptions>=1.10.0",
+    "robot_descriptions>=1.18.0",
     "gdown>=4.6.6",
     "plyfile",
     "opencv-python",

src/viser/client/src/ControlPanel/Generated.tsx

@@ -79,16 +79,23 @@ function GuiContainer({ containerUuid }: { containerUuid: string }) {
   );
   const out = (
     <Box pt="xs">
-      {guiUuidOrderPairArray.map((pair) => (
-        <GeneratedInput key={pair.uuid} guiUuid={pair.uuid} />
+      {guiUuidOrderPairArray.map((pair, index) => (
+        <GeneratedInput
+          key={pair.uuid}
+          guiUuid={pair.uuid}
+          nextGuiUuid={guiUuidOrderPairArray[index + 1]?.uuid ?? null}
+        />
       ))}
     </Box>
   );
   return out;
 }
 
 /** A single generated GUI element. */
-function GeneratedInput(props: { guiUuid: string }) {
+function GeneratedInput(props: {
+  guiUuid: string;
+  nextGuiUuid: string | null;
+}) {
   const viewer = React.useContext(ViewerContext)!;
   const conf = viewer.useGui((state) => state.guiConfigFromUuid[props.guiUuid]);
   if (conf === undefined) {
@@ -97,7 +104,7 @@ function GeneratedInput(props: { guiUuid: string }) {
   }
   switch (conf.type) {
     case "GuiFolderMessage":
-      return <FolderComponent {...conf} />;
+      return <FolderComponent {...conf} nextGuiUuid={props.nextGuiUuid} />;
     case "GuiTabGroupMessage":
       return <TabGroupComponent {...conf} />;
     case "GuiMarkdownMessage":

src/viser/client/src/components/Folder.css.ts

@@ -6,14 +6,8 @@ export const folderWrapper = style({
   position: "relative",
   marginLeft: vars.spacing.xs,
   marginRight: vars.spacing.xs,
-  // If there's a GUI element above, we need more margin.
   marginTop: vars.spacing.xs,
-  // If there's a GUI element below, we need more margin.
-  // Note: 0.5em is the vertical margin below general GUI elements.
-  marginBottom: "1.2em",
-  ":last-child": {
-    marginBottom: "0.5em",
-  },
+  marginBottom: vars.spacing.xs,
   paddingBottom: `calc(${vars.spacing.xs} - 0.5em)`,
 });
 

src/viser/client/src/components/Folder.tsx

@@ -10,19 +10,27 @@ import { folderLabel, folderToggleIcon, folderWrapper } from "./Folder.css";
 export default function FolderComponent({
   uuid,
   props: { label, visible, expand_by_default },
-}: GuiFolderMessage) {
+  nextGuiUuid,
+}: GuiFolderMessage & { nextGuiUuid: string | null }) {
   const viewer = React.useContext(ViewerContext)!;
   const [opened, { toggle }] = useDisclosure(expand_by_default);
   const guiIdSet = viewer.useGui(
     (state) => state.guiUuidSetFromContainerUuid[uuid],
   );
   const guiContext = React.useContext(GuiComponentContext)!;
   const isEmpty = guiIdSet === undefined || Object.keys(guiIdSet).length === 0;
+  const nextGuiType = viewer.useGui((state) =>
+    nextGuiUuid == null ? null : state.guiConfigFromUuid[nextGuiUuid]?.type,
+  );
 
   const ToggleIcon = opened ? IconChevronUp : IconChevronDown;
   if (!visible) return null;
   return (
-    <Paper withBorder className={folderWrapper}>
+    <Paper
+      withBorder
+      className={folderWrapper}
+      mb={nextGuiType === "GuiFolderMessage" ? "md" : undefined}
+    >
       <Paper
         className={folderLabel}
         style={{