Add a .matrix property to mujoco.Camera that exposes the camera matrix

yuvaltassa · alimuldal · commit 45fb1e668fc2 · 2020-09-25T16:03:12.000+01:00
PiperOrigin-RevId: 330349436
Change-Id: Ib5295f0e3501b5d5c39e45c3eb2b09822322c311
diff --git a/dm_control/mujoco/engine.py b/dm_control/mujoco/engine.py
@@ -686,6 +686,46 @@ def scene(self):
     """Returns the `mujoco.MjvScene` instance used by the camera."""
     return self._scene
 
+  @property
+  def matrix(self):
+    """Returns the 3x4 camera matrix.
+
+       For a description of the camera matrix see, e.g.,
+       https://en.wikipedia.org/wiki/Camera_matrix.
+       For a usage example, see the associated test.
+    """
+    camera_id = self._render_camera.fixedcamid
+    if camera_id == -1:
+      # If the camera is a 'free' camera, we get its position and orientation
+      # from the scene data structure. It is a stereo camera, so we average over
+      # the left and right channels. Note: we call `self.update()` in order to
+      # ensure that the contents of `scene.camera` are correct.
+      self.update()
+      pos = np.mean(self.scene.camera.pos, axis=0)
+      z = -np.mean(self.scene.camera.forward, axis=0)
+      y = np.mean(self.scene.camera.up, axis=0)
+      rot = np.vstack((np.cross(y, z), y, z))
+      fov = self._physics.model.vis.global_.fovy
+    else:
+      pos = self._physics.data.cam_xpos[camera_id]
+      rot = self._physics.data.cam_xmat[camera_id].reshape(3, 3).T
+      fov = self._physics.model.cam_fovy[camera_id]
+
+    # Translation matrix (4x4).
+    translation = np.eye(4)
+    translation[0:3, 3] = -pos
+    # Rotation matrix (4x4).
+    rotation = np.eye(4)
+    rotation[0:3, 0:3] = rot
+    # Focal transformation matrix (3x4).
+    focal_scaling = (1./np.tan(np.deg2rad(fov)/2)) * self.height / 2.0
+    focal = np.diag([-focal_scaling, focal_scaling, 1.0, 0])[0:3, :]
+    # Image matrix (3x3).
+    image = np.eye(3)
+    image[0, 2] = (self.width - 1) / 2.0
+    image[1, 2] = (self.height - 1) / 2.0
+    return image.dot(focal).dot(rotation).dot(translation)
+
   def update(self, scene_option=None):
     """Updates geometry used for rendering.
 
diff --git a/dm_control/mujoco/engine_test.py b/dm_control/mujoco/engine_test.py
@@ -191,6 +191,60 @@ def testCameraSelection(self, coordinates, expected_selection):
     selected = camera.select(coordinates)
     self.assertEqual(expected_selection, selected[:2])
 
+  @parameterized.parameters(
+      dict(camera_id='cam0', height=200, width=300),
+      dict(camera_id=1, height=300, width=200),
+      dict(camera_id=-1, height=400, width=400),
+  )
+  def testCameraMatrix(self, camera_id, height, width):
+    """Tests the camera_matrix() method.
+
+       Creates a model with two cameras and two small geoms. We render the scene
+       with one of the cameras and check that the geom locations, projected into
+       pixel space, are correct, using segmenation rendering.
+       xyz2pixels() shows how the transformation is used. For a description
+       of the camera matrix see https://en.wikipedia.org/wiki/Camera_matrix.
+
+    Args:
+      camera_id: One of the two cameras. Can be either integer or String.
+      height: The height of the image (pixels).
+      width: The width of the image (pixels).
+    """
+
+    def xyz2pixels(x, y, z, camera_matrix):
+      """Transforms from world coordinates to pixel coordinates."""
+      xs, ys, s = camera_matrix.dot(np.array([x, y, z, 1.0]))
+      return xs/s, ys/s
+
+    two_geoms_and_two_cameras = """
+    <mujoco>
+      <visual>
+        <global fovy="55"/>
+      </visual>
+      <worldbody>
+        <light name="top" pos="0 0 1"/>
+        <geom name="red" pos=".2 0 0" size=".005" rgba="1 0 0 1"/>
+        <geom name="green" pos=".2 .2 .1" size=".005" rgba="0 1 0 1"/>
+        <camera name="cam0" pos="1 .5 1" zaxis="1 .5 1" fovy="20"/>
+        <camera name="cam1" pos=".1 .1 1" xyaxes="1 1 0 -1 0 0"/>
+      </worldbody>
+    </mujoco>
+    """
+    physics = engine.Physics.from_xml_string(two_geoms_and_two_cameras)
+    camera = engine.Camera(physics, width=width, height=height,
+                           camera_id=camera_id)
+    camera_matrix = camera.matrix  # Get camera matrix.
+    pixels = camera.render(segmentation=True)  # Render a segmentation frame.
+    for geom_id in [0, 1]:
+      # Compute the location of the geom in pixel space using the camera matrix.
+      x, y = xyz2pixels(*physics.data.geom_xpos[geom_id], camera_matrix)
+      row = int(round(y))
+      column = int(round(x))
+      # Compare segmentation values of nearest pixel to corresponding geom.
+      [obj_id, obj_type] = pixels[row, column, :]
+      self.assertEqual(obj_type, enums.mjtObj.mjOBJ_GEOM)
+      self.assertEqual(obj_id, geom_id)
+
   def testMovableCameraSetGetPose(self):
     height, width = 240, 320
 
diff --git a/setup.py b/setup.py
@@ -177,7 +177,7 @@ def is_excluded(s):
 
 setup(
     name='dm_control',
-    version='0.0.329943179',
+    version='0.0.330349436',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
