imt-ring 1.2.1__py3-none-any.whl

ring/rendering/vispy_render.py

@@ -0,0 +1,340 @@
+from abc import ABC
+from abc import abstractmethod
+from abc import abstractstaticmethod
+from functools import partial
+from typing import Optional, TypeVar
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from ring import algebra
+from ring import base
+from ring import maths
+from tree_utils import PyTree
+from tree_utils import tree_batch
+from vispy import scene
+from vispy.scene import MatrixTransform
+
+from . import vispy_visuals
+
+Camera = TypeVar("Camera")
+Visual = TypeVar("Visual")
+VisualPosOri1 = PyTree
+VisualPosOri2 = PyTree
+
+
+class Scene(ABC):
+    _xyz: bool = True
+    _xyz_root: bool = True
+    _xyz_transform1: bool = True
+    visuals: list[Visual] = []
+
+    """
+    Example:
+        >> renderer = Renderer()
+        >> renderer.init(sys.geoms)
+        >> for t in range(xs.shape()):
+        >>   renderer.update(xs[t])
+        >>   image = renderer.render()
+    """
+
+    @abstractmethod
+    def _get_camera(self) -> Camera:
+        pass
+
+    @abstractmethod
+    def _set_camera(self, camera: Camera) -> None:
+        pass
+
+    @abstractmethod
+    def _render(self) -> jax.Array:
+        pass
+
+    def enable_xyz(self, enable_root: bool = True) -> None:
+        self._xyz = True
+        if enable_root:
+            self._xyz_root = True
+
+    def disable_xyz(self, disable_root: bool = True) -> None:
+        self._xyz = False
+        if disable_root:
+            self._xyz_root = False
+
+    def enable_xyz_transform1(self):
+        self._xyz_transform1 = True
+
+    def disable_xyz_tranform1(self):
+        self._xyz_transform1 = False
+
+    def render(
+        self, camera: Optional[Camera | list[Camera]] = None
+    ) -> jax.Array | list[jax.Array]:
+        "Returns: RGBA Array of Shape = (M, N, 4)"
+        if camera is None:
+            camera = self._get_camera()
+
+        if not isinstance(camera, list):
+            self._set_camera(camera)
+            return self._render()
+
+        images = []
+        for cam in camera:
+            self._set_camera(cam)
+            images.append(self._render())
+        return images
+
+    def _add_box(self, box: base.Box) -> Visual:
+        raise NotImplementedError
+
+    def _add_sphere(self, sphere: base.Sphere) -> Visual:
+        raise NotImplementedError
+
+    def _add_cylinder(self, cyl: base.Cylinder) -> Visual:
+        raise NotImplementedError
+
+    def _add_capsule(self, cap: base.Capsule) -> Visual:
+        raise NotImplementedError
+
+    def _add_xyz(self) -> Visual:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _remove_visual(self, visual: Visual) -> None:
+        pass
+
+    def _remove_all_visuals(self):
+        for visual in self.visuals:
+            self._remove_visual(visual)
+
+    def init(self, geoms: list[base.Geometry]):
+        self._remove_all_visuals()
+
+        self.geoms = [geom for geom in geoms]
+        self._fresh_init = True
+
+        geom_link_idx = []
+        geom_transform = []
+        self.visuals = []
+        for geom in geoms:
+            geom_link_idx.append(geom.link_idx)
+            geom_transform.append(geom.transform)
+            if isinstance(geom, base.Box):
+                visual = self._add_box(geom)
+            elif isinstance(geom, base.Sphere):
+                visual = self._add_sphere(geom)
+            elif isinstance(geom, base.Cylinder):
+                visual = self._add_cylinder(geom)
+            elif isinstance(geom, base.Capsule):
+                visual = self._add_capsule(geom)
+            elif isinstance(geom, base.XYZ):
+                visual = self._add_xyz()
+                if not self._xyz_transform1:
+                    geom_transform.pop()
+                    geom_transform.append(base.Transform.zero())
+            else:
+                raise Exception(f"Unknown geom type: {type(geom)}")
+            self.visuals.append(visual)
+
+        if self._xyz:
+            unique_link_indices = np.unique(np.array(geom_link_idx))
+            for unique_link_idx in unique_link_indices:
+                geom_link_idx.append(unique_link_idx)
+                geom_transform.append(base.Transform.zero())
+                self.visuals.append(self._add_xyz())
+                # otherwise the .update function won't iterate
+                # over all visuals since it uses a zip(...)
+                self.geoms.append(None)
+
+        if self._xyz_root:
+            # add one final for root frame
+            self._add_xyz()
+
+        self.geom_link_idx = tree_batch(geom_link_idx, backend="jax")
+        self.geom_transform = tree_batch(geom_transform, backend="jax")
+
+    @abstractstaticmethod
+    def _compute_transform_per_visual(
+        x_links: base.Transform,
+        x_link_to_geom: base.Transform,
+        geom_link_idx: int,
+    ) -> VisualPosOri1:
+        "This can easily account for possible convention differences"
+        pass
+
+    @abstractstaticmethod
+    def _postprocess_transforms(transform: VisualPosOri1) -> VisualPosOri2:
+        pass
+
+    @abstractmethod
+    def _init_visual(
+        self, visual: Visual, transform: VisualPosOri2, geom: None | base.Geometry
+    ):
+        pass
+
+    def _update_visual(
+        self, visual: Visual, transform: VisualPosOri2, geom: None | base.Geometry
+    ):
+        self._init_visual(visual, transform, geom)
+
+    def update(self, x: base.Transform):
+        "`x` are (n_links,) Transforms."
+
+        # step 1: pre-compute all required transforms
+        transform_per_visual = _compile_staticmethod(
+            self._compute_transform_per_visual,
+            x,
+            self.geom_transform,
+            self.geom_link_idx,
+        )
+
+        # step 2: postprocess all transforms once
+        transform_per_visual = self._postprocess_transforms(transform_per_visual)
+
+        # step 3: update visuals
+        for i, (visual, geom) in enumerate(zip(self.visuals, self.geoms)):
+            t = jax.tree_map(lambda arr: arr[i], transform_per_visual)
+            if self._fresh_init:
+                self._init_visual(visual, t, geom)
+            else:
+                self._update_visual(visual, t, geom)
+
+        # step 4: unset flag
+        self._fresh_init = False
+
+
+@partial(jax.jit, static_argnums=0)
+def _compile_staticmethod(static_method, x, geom_transform, geom_link_idx):
+    return jax.vmap(static_method, in_axes=(None, 0, 0))(
+        x, geom_transform, geom_link_idx
+    )
+
+
+class VispyScene(Scene):
+    def __init__(
+        self,
+        show_cs=False,
+        show_cs_root=True,
+        width: int = 320,
+        height: int = 240,
+        camera: scene.cameras.BaseCamera = scene.TurntableCamera(
+            elevation=25, distance=4.0, azimuth=25
+        ),
+        **kwargs,
+    ):
+        """Scene which can be rendered.
+
+        Args:
+            geoms (list[list[Geometry]]): A list of list of geometries per link.
+                len(geoms) == number of links in system
+            show_cs (bool, optional): Show coordinate system of links.
+                Defaults to True.
+            show_cs_root (bool, optional): Show coordinate system of earth frame.
+                Defaults to True.
+            camera (scene.cameras.BaseCamera, optional): The camera angle.
+                Defaults to scene.TurntableCamera( elevation=30, distance=6 ).
+
+        Example:
+            >> scene = VispyScene()
+            >> scene.init(sys.geoms)
+            >> scene.update(state.x)
+            >> image = scene.render()
+        """
+        self.canvas = scene.SceneCanvas(
+            keys="interactive", size=(width, height), show=True, **kwargs
+        )
+        self.view = self.canvas.central_widget.add_view()
+        self._set_camera(camera)
+        if show_cs:
+            self.enable_xyz()
+        else:
+            self.disable_xyz(not show_cs_root)
+
+    def _set_camera(self, camera: scene.cameras.BaseCamera) -> None:
+        self.view.camera = camera
+
+    def _get_camera(self) -> scene.cameras.BaseCamera:
+        return self.view.camera
+
+    def _render(self) -> jax.Array:
+        return self.canvas.render(alpha=True)
+
+    def _add_box(self, box: base.Box) -> Visual:
+        return vispy_visuals.Box(
+            box.dim_x,
+            box.dim_z,
+            box.dim_y,
+            color=box.color,
+            edge_color=box.edge_color,
+            parent=self.view.scene,
+        )
+
+    def _add_sphere(self, sphere: base.Sphere) -> Visual:
+        return vispy_visuals.Sphere(
+            sphere.radius,
+            color=sphere.color,
+            edge_color=sphere.edge_color,
+            parent=self.view.scene,
+        )
+
+    def _add_cylinder(self, cyl: base.Cylinder) -> Visual:
+        return vispy_visuals.Cylinder(
+            cyl.radius,
+            cyl.length,
+            color=cyl.color,
+            edge_color=cyl.edge_color,
+            parent=self.view.scene,
+        )
+
+    def _add_capsule(self, cap: base.Capsule) -> Visual:
+        return vispy_visuals.Capsule(
+            cap.radius,
+            cap.length,
+            color=cap.color,
+            edge_color=cap.edge_color,
+            parent=self.view.scene,
+        )
+
+    def _add_xyz(self) -> Visual:
+        return scene.visuals.XYZAxis(parent=self.view.scene)
+
+    def _remove_visual(self, visual: scene.visuals.VisualNode) -> None:
+        visual.parent = None
+
+    @staticmethod
+    def _compute_transform_per_visual(
+        x_links: base.Transform,
+        x_link_to_geom: base.Transform,
+        geom_link_idx: int,
+    ) -> jax.Array:
+        x = jax.lax.cond(
+            geom_link_idx == -1,
+            lambda: base.Transform.zero(),
+            lambda: x_links[geom_link_idx],
+        )
+        x = algebra.transform_mul(x_link_to_geom, x)
+        E = maths.quat_to_3x3(x.rot)
+        M = jnp.eye(4)
+        M = M.at[:3, :3].set(E)
+        T = jnp.eye(4)
+        T = T.at[3, :3].set(x.pos)
+        return M @ T
+
+    @staticmethod
+    def _postprocess_transforms(transform: jax.Array) -> np.ndarray:
+        return np.asarray(transform)
+
+    def _init_visual(
+        self,
+        visual: scene.visuals.VisualNode,
+        transform: np.ndarray,
+        geom: base.Geometry,
+    ):
+        visual.transform = MatrixTransform(transform)
+
+    def _update_visual(
+        self,
+        visual: scene.visuals.VisualNode,
+        transform: np.ndarray,
+        geom: base.Geometry,
+    ):
+        visual.transform.matrix = transform

ring/rendering/vispy_visuals.py

@@ -0,0 +1,290 @@
+import numpy as np
+from ring.base import Color
+from vispy.geometry.meshdata import MeshData
+from vispy.scene.visuals import create_visual_node
+from vispy.scene.visuals import Mesh
+from vispy.visuals import CompoundVisual
+from vispy.visuals import SphereVisual as _SphereVisual
+from vispy.visuals import TubeVisual
+
+# vertex density per unit length
+_vectices_per_unit_length = 10
+
+_default_color = (1, 0.8, 0.7, 1)
+_default_edge_color = "black"
+
+
+class DoubleMeshVisual(CompoundVisual):
+    _lines: Mesh
+    _faces: Mesh
+
+    def __init__(
+        self, verts, edges, faces, *, color: Color = None, edge_color: Color = None
+    ):
+        if color is None and edge_color is None:
+            color = _default_color
+
+        if color is not None:
+            self._faces = Mesh(verts, faces, color=color, shading=None)
+            self.light_dir = np.array([0, -1, 0])
+        else:
+            self._faces = Mesh()
+
+        if edge_color is not None:
+            self._edges = Mesh(verts, edges, color=edge_color, mode="lines")
+        else:
+            self._edges = Mesh()
+
+        super().__init__([self._faces, self._edges])
+        self._faces.set_gl_state(
+            polygon_offset_fill=True, polygon_offset=(1, 1), depth_test=True
+        )
+
+
+class SphereVisual(_SphereVisual):
+    def __init__(self, radius: float, color: Color = None, edge_color: Color = None):
+        if color is None and edge_color is None:
+            color = _default_color
+
+        radius = float(radius)
+
+        num_rows = max(int(np.pi * radius * _vectices_per_unit_length), 10)
+        num_cols = max(int(2 * np.pi * radius * _vectices_per_unit_length), 20)
+
+        super().__init__(
+            radius,
+            color=color,
+            edge_color=edge_color,
+            rows=num_rows,
+            cols=num_cols,
+            method="latitude",
+            shading="smooth",
+        )
+
+
+Sphere = create_visual_node(SphereVisual)
+
+
+def box_mesh(
+    dim_x: float, dim_y: float, dim_z: float
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    verts = np.array(
+        [
+            (-dim_x, -dim_y, -dim_z),
+            (dim_x, -dim_y, -dim_z),
+            (-dim_x, dim_y, -dim_z),
+            (dim_x, dim_y, -dim_z),
+            (-dim_x, -dim_y, dim_z),
+            (dim_x, -dim_y, dim_z),
+            (-dim_x, dim_y, dim_z),
+            (dim_x, dim_y, dim_z),
+        ],
+        dtype=np.float32,
+    )
+
+    verts /= 2
+
+    edges = np.array(
+        [
+            (0, 1),
+            (0, 2),
+            (0, 4),
+            (1, 3),
+            (1, 5),
+            (2, 3),
+            (2, 6),
+            (3, 7),
+            (4, 5),
+            (4, 6),
+            (5, 7),
+            (6, 7),
+        ],
+        dtype=np.uint32,
+    )
+
+    faces = np.array(
+        [
+            (0, 1, 2),
+            (1, 2, 3),
+            (0, 1, 4),
+            (1, 4, 5),
+            (0, 2, 4),
+            (2, 4, 6),
+            (1, 3, 5),
+            (3, 5, 7),
+            (2, 3, 6),
+            (3, 6, 7),
+            (4, 5, 6),
+            (5, 6, 7),
+        ]
+    )
+
+    return verts, edges, faces
+
+
+class BoxVisual(DoubleMeshVisual):
+    # NOTE: need a custom BoxVisual class, since vispy.scene.visuals.Box does not
+    # support shading
+
+    def __init__(
+        self,
+        dim_x: float,
+        dim_y: float,
+        dim_z: float,
+        *,
+        color: Color = None,
+        edge_color: Color = None
+    ):
+        if color is None:
+            color = _default_color
+
+        if edge_color is None:
+            edge_color = _default_edge_color
+
+        dim_x = float(dim_x)
+        dim_y = float(dim_y)
+        dim_z = float(dim_z)
+
+        self.dim_x = dim_x
+        self.dim_y = dim_y
+        self.dim_z = dim_z
+
+        verts, edges, faces = box_mesh(dim_x, dim_y, dim_z)
+
+        super().__init__(verts, edges, faces, color=color, edge_color=edge_color)
+
+
+Box = create_visual_node(BoxVisual)
+
+
+class CylinderVisual(TubeVisual):
+    def __init__(
+        self,
+        radius: float,
+        length: float,
+        *,
+        color: Color = None,
+        edge_color: Color = None
+    ):
+        if color is None and edge_color is None:
+            color = _default_color
+
+        radius = float(radius)
+        length = float(length)
+
+        num_length_points = 10 * max(int(length * _vectices_per_unit_length), 10)
+        num_radial_points = max(int(2 * np.pi * radius * _vectices_per_unit_length), 20)
+
+        points = np.zeros((num_length_points, 3))
+        points[:, 0] = np.linspace(-length / 2, length / 2, num_length_points)
+
+        self.radius = radius
+        self.length = length
+
+        super().__init__(
+            points,
+            radius,
+            tube_points=num_radial_points,
+            closed=True,
+            color=color,
+            shading="smooth",
+        )
+
+
+Cylinder = create_visual_node(CylinderVisual)
+
+
+def capsule_mesh(radius: float, length: float, offset: bool = True) -> MeshData:
+    if length < 2 * radius:
+        raise ValueError("length must be at least 2 * radius")
+
+    # number of cap vertices in x direction
+    num_sphere_rows = max(int(radius * _vectices_per_unit_length), 10)
+
+    # length without caps
+    cyl_length = length - 2 * radius
+    # number of cylinder vertices in x direction
+    num_cyl_rows = max(int(cyl_length * _vectices_per_unit_length), 10)
+
+    num_total_rows = 2 * num_sphere_rows + num_cyl_rows
+
+    # number of radial vertices
+    num_cols = max(int(2 * np.pi * radius * _vectices_per_unit_length), 20)
+
+    verts = np.empty((num_total_rows, num_cols, 3), dtype=np.float32)
+
+    # polar angle
+    theta_top = np.linspace(0.0, np.pi / 2, num_sphere_rows)
+    theta_bottom = np.linspace(np.pi / 2, np.pi, num_sphere_rows)
+
+    # fill in x coordinate
+    verts[:num_sphere_rows, :, 0] = radius * np.cos(theta_top[:, None]) + cyl_length / 2
+
+    verts[num_sphere_rows:-num_sphere_rows, :, 0] = np.linspace(
+        -cyl_length / 2, cyl_length / 2, num_cyl_rows
+    )[::-1, None]
+
+    verts[-num_sphere_rows:, :, 0] = (
+        radius * np.cos(theta_bottom[:, None]) - cyl_length / 2
+    )
+
+    # azimuth angle
+    phi = (np.linspace(0, 2 * np.pi, num_cols))[None, :]
+
+    if offset:
+        # rotate each row by 1/2 column
+        phi = phi + (np.pi / num_cols) * np.arange(num_total_rows)[:, None]
+
+    # y and z coordinates
+    verts[..., 1] = radius * np.cos(phi)
+    verts[..., 2] = radius * np.sin(phi)
+
+    # for caps: bend inwards to close
+    verts[:num_sphere_rows, :, 1:3] *= np.sin(theta_top[:, None, None])
+    verts[-num_sphere_rows:, :, 1:3] *= np.sin(theta_bottom[:, None, None])
+
+    verts = verts.reshape(-1, 3)
+
+    # compute faces
+    faces = np.empty(((num_total_rows - 1) * num_cols * 2, 3), dtype=np.uint32)
+
+    rowtemplate1 = (
+        (np.arange(num_cols).reshape(num_cols, 1) + np.array([[0, 1, 0]])) % num_cols
+    ) + np.array([[num_cols, 0, 0]])
+
+    rowtemplate2 = (
+        (np.arange(num_cols).reshape(num_cols, 1) + np.array([[1, 1, 0]])) % num_cols
+    ) + np.array([[num_cols, 0, num_cols]])
+
+    for row in range(num_total_rows - 1):
+        start = row * num_cols * 2
+
+        faces[start : start + num_cols] = rowtemplate1 + row * num_cols
+        faces[start + num_cols : start + (num_cols * 2)] = rowtemplate2 + row * num_cols
+
+    mesh = MeshData(vertices=verts, faces=faces)
+
+    return mesh.get_vertices(), mesh.get_edges(), mesh.get_faces()
+
+
+class CapsuleVisual(DoubleMeshVisual):
+    def __init__(
+        self,
+        radius: float,
+        length: float,
+        *,
+        color: Color = None,
+        edge_color: Color = None
+    ):
+        radius = float(radius)
+        length = float(length)
+
+        self.radius = radius
+        self.length = length
+
+        verts, edges, faces = capsule_mesh(radius, length)
+
+        super().__init__(verts, edges, faces, color=color, edge_color=edge_color)
+
+
+Capsule = create_visual_node(CapsuleVisual)

ring/sim2real/__init__.py

@@ -0,0 +1,7 @@
+from .sim2real import delete_to_world_pos_rot
+from .sim2real import match_xs
+from .sim2real import randomize_to_world_pos_rot
+from .sim2real import scale_xs
+from .sim2real import unzip_xs
+from .sim2real import xs_from_raw
+from .sim2real import zip_xs