Merge pull request #618 from Genesis-Embodied-AI/yian/mounted_camera

Yian/mounted camera

Author: wangyian-me

genesis/vis/camera.py

@@ -89,6 +89,8 @@ def __init__(
         self._aspect_ratio = self._res[0] / self._res[1]
         self._visualizer = visualizer
         self._is_built = False
+        self._attached_link = None
+        self._attached_offset_T = None
 
         self._in_recording = False
         self._recorded_imgs = []
@@ -111,6 +113,27 @@ def _build(self):
         self._is_built = True
         self.set_pose(self._transform, self._pos, self._lookat, self._up)
 
+    def attach(self, rigid_link, offset_T):
+        self._attached_link = rigid_link
+        self._attached_offset_T = offset_T
+
+    def detach(self):
+        self._attached_link = None
+        self._attached_offset_T = None
+
+    @gs.assert_built
+    def move_to_attach(self):
+        if self._attached_link is None:
+            gs.raise_exception(f"The camera hasn't been mounted!")
+        if self._visualizer._scene.n_envs > 0:
+            gs.raise_exception(f"Mounted camera not supported in parallel simulation!")
+
+        link_pos = self._attached_link.get_pos().cpu().numpy()
+        link_quat = self._attached_link.get_quat().cpu().numpy()
+        link_T = gu.trans_quat_to_T(link_pos, link_quat)
+        transform = link_T @ self._attached_offset_T
+        self.set_pose(transform=transform)
+
     @gs.assert_built
     def render(self, rgb=True, depth=False, segmentation=False, colorize_seg=False, normal=False):
         """

genesis/vis/visualizer.py

@@ -144,6 +144,10 @@ def update_visual_states(self):
         if self._t < self._scene._t:
             self._t = self._scene._t
 
+            for camera in self._cameras:
+                if camera._attached_link is not None:
+                    camera.move_to_attach()
+
             if self._scene.rigid_solver.is_active():
                 self._scene.rigid_solver.update_geoms_render_T()
                 self._scene.rigid_solver._kernel_update_vgeoms()