feat(viser): add viser viewer (#628)

Signed-off-by: Hirokazu Ishida (SB Intuitions) <[email protected]>

Author: HiroIshida

requirements.txt

@@ -26,3 +26,4 @@ scipy<=1.2.3;python_version<"3.0"
 scipy>=1.6.3;python_version>="3.8"
 six
 trimesh>=3.9.0,!=3.23.0
+viser

skrobot/model/primitives.py

@@ -101,6 +101,16 @@ def __init__(self,
 
         super(Axis, self).__init__(pos=pos, rot=rot, name=name,
                                    visual_mesh=visual_mesh)
+        self._axis_length = axis_length
+        self._axis_radius = axis_radius
+
+    @property
+    def axis_length(self):
+        return self._axis_length
+
+    @property
+    def axis_radius(self):
+        return self._axis_radius
 
     @classmethod
     def from_coords(cls, coords, **kwargs):
@@ -136,15 +146,18 @@ def __init__(self, extents, vertex_colors=None, face_colors=None,
         super(Box, self).__init__(pos=pos, rot=rot, name=name,
                                   collision_mesh=mesh,
                                   visual_mesh=mesh)
-        self.extents = extents
-        self._extents = extents  # for backward compatibility
+        self._extents = extents
         if with_sdf:
             sdf = BoxSDF(extents)
             self.assoc(sdf, relative_coords="local")
             self._sdf = sdf
         else:
             self._sdf = None
 
+    @property
+    def extents(self):
+        return self._extents
+
 
 class CameraMarker(Link):
 
@@ -244,14 +257,18 @@ def __init__(self, radius, subdivisions=3, color=None,
                                      collision_mesh=mesh,
                                      visual_mesh=mesh)
 
-        self.radius = radius
+        self._radius = radius
         if with_sdf:
             sdf = SphereSDF(radius)
             self.assoc(sdf, relative_coords="local")
             self._sdf = sdf
         else:
             self._sdf = None
 
+    @property
+    def radius(self):
+        return self._radius
+
 
 class Annulus(Link):
 

skrobot/viewers/__init__.py

@@ -66,3 +66,9 @@ class PyrenderViewer(DummyViewer):
 except ImportError:
     class JupyterNotebookViewer(DummyViewer):
         pass
+
+try:
+    from ._viser import ViserVisualizer
+except ImportError:
+    class ViserVisualizer(DummyViewer):
+        pass

skrobot/viewers/_viser.py

@@ -0,0 +1,124 @@
+from typing import Union
+import webbrowser
+
+import numpy as np
+import trimesh
+import viser
+
+from skrobot.coordinates.math import matrix2quaternion
+from skrobot.model.link import Link
+from skrobot.model.primitives import Axis
+from skrobot.model.primitives import LineString
+from skrobot.model.primitives import PointCloudLink
+from skrobot.model.primitives import Sphere
+from skrobot.model.robot_model import CascadedLink
+
+
+class ViserVisualizer:
+    def __init__(self, draw_grid: bool = True):
+        self._server = viser.ViserServer()
+        self._linkid_to_handle = dict()
+        self._linkid_to_link = dict()
+
+    def draw_grid(self, width: float = 20.0, height: float = -0.001):
+        self._server.scene.add_grid(
+            "/grid",
+            width=20.0,
+            height=20.0,
+            position=np.array([0.0, 0.0, -0.01]),
+        )
+
+    def _add_link(self, link: Link):
+        assert isinstance(link, Link)
+        link_id = str(id(link))
+        if link_id in self._linkid_to_handle:
+            return
+
+        handle = None
+        if isinstance(link, Sphere):
+            # Although sphere can be treated as trimesh, naively rendering
+            # it requires high cost. Therefore, we use an analytic sphere.
+            color = link.visual_mesh.visual.face_colors[0, :3]
+            alpha = link.visual_mesh.visual.face_colors[0, 3]
+            if alpha > 1.0:
+                alpha = alpha / 255.0
+            handle = self._server.scene.add_icosphere(
+                link.name,
+                radius=link.radius,
+                position=link.worldpos(),
+                color=color,
+                opacity=alpha)
+        elif isinstance(link, Axis):
+            handle = self._server.scene.add_frame(
+                    link.name,
+                    axes_length=link.axis_length,
+                    axes_radius=link.axis_radius,
+                    wxyz=matrix2quaternion(link.worldrot()),
+                    position=link.worldpos(),
+                )
+        elif isinstance(link, PointCloudLink):
+            mesh = link.visual_mesh
+            assert isinstance(mesh, trimesh.PointCloud)
+            if len(mesh.colors) > 0:
+                colors = mesh.colors[:, :3]
+            else:
+                colors = np.zeros(3)
+            self._server.scene.add_point_cloud(
+                    link.name,
+                    points=mesh.vertices,
+                    colors=colors,
+                    point_size=0.002,  # TODO(HiroIshida): configurable
+                )
+        elif isinstance(link, LineString):
+            raise NotImplementedError("not implemented yet")
+        else:
+            mesh = link.concatenated_visual_mesh
+            if mesh is not None:
+                handle = self._server.scene.add_mesh_trimesh(
+                        link.name,
+                        mesh=mesh,
+                        wxyz=matrix2quaternion(link.worldrot()),
+                        position=link.worldpos(),
+                    )
+
+        if handle is not None:
+            self._linkid_to_link[link_id] = link
+            self._linkid_to_handle[link_id] = handle
+
+    def add(self, geometry: Union[Link, CascadedLink]):
+        if isinstance(geometry, Link):
+            self._add_link(geometry)
+        elif isinstance(geometry, CascadedLink):
+            for link in geometry.link_list:
+                self._add_link(link)
+        else:
+            raise TypeError("geometry must be Link or CascadedLink")
+
+    def show(self):
+        host = self._server.get_host()
+        port = self._server.get_port()
+        url = f"http://{host}:{port}"
+        webbrowser.open(url)
+
+    def redraw(self):
+        for link_id, handle in self._linkid_to_handle.items():
+            link = self._linkid_to_link[link_id]
+            handle.position = link.worldpos()
+            handle.wxyz = matrix2quaternion(link.worldrot())
+
+    def delete(self, geometry: Union[Link, CascadedLink]):
+        if isinstance(geometry, Link):
+            links = [geometry]
+        elif isinstance(geometry, CascadedLink):
+            links = geometry.link_list
+        else:
+            raise TypeError("geometry must be Link or CascadedLink")
+
+        for link in links:
+            link_id = str(id(link))
+            if link_id not in self._linkid_to_handle:
+                continue
+            handle = self._linkid_to_handle[link_id]
+            handle.remove()
+            self._linkid_to_link.pop(link_id)
+            self._linkid_to_handle.pop(link_id)