Add get/set for cameras

mojo/elements/camera.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING
 
+import mujoco
 import numpy as np
 from mujoco_utils import mjcf_utils
 from typing_extensions import Self
@@ -59,7 +60,20 @@ def set_position(self, position: np.ndarray):
         self.mjcf.pos = position
 
     def get_position(self) -> np.ndarray:
-        return self._mojo.physics.bind(self.mjcf).pos
+        return self._mojo.physics.bind(self.mjcf).pos.copy()
+
+    def set_quaternion(self, quaternion: np.ndarray):
+        # wxyz
+        quaternion = np.array(quaternion)  # ensure is numpy array
+        mat = np.zeros(9)
+        mujoco.mju_quat2Mat(mat, quaternion)
+        self._mojo.physics.bind(self.mjcf).xmat = mat
+        self.mjcf.quat = quaternion
+
+    def get_quaternion(self) -> np.ndarray:
+        quat = np.zeros(4)
+        mujoco.mju_mat2Quat(quat, self._mojo.physics.bind(self.mjcf).xmat)
+        return quat
 
     def set_focal(self, focal: np.ndarray):
         if self.mjcf.sensorsize is None:
@@ -73,16 +87,29 @@ def get_focal(self) -> np.ndarray:
         return self.mjcf.focal
 
     def set_sensor_size(self, sensor_size: np.ndarray):
-        if self.mjcf.focal is None:
+        # Either focal or focalpixel must be set
+        if self.mjcf.focal is None or self.mjcf.focalpixel is None:
             self.mjcf.focal = np.array([0, 0])
+        # Resolution must be set
         if self.mjcf.resolution is None:
             self.mjcf.resolution = np.array([1, 1])
-        self.mjcf.sensorsize = sensor_size
+        self.mjcf.sensorsize = np.array(sensor_size)
         self._mojo.mark_dirty()
 
     def get_sensor_size(self) -> np.ndarray:
         return self.mjcf.sensorsize
 
+    def set_focal_pixel(self, focal_pixel: np.ndarray):
+        if self.mjcf.sensorsize is None:
+            self.mjcf.sensorsize = np.array([0, 0])
+        if self.mjcf.resolution is None:
+            self.mjcf.resolution = np.array([1, 1])
+        self.mjcf.focalpixel = focal_pixel
+        self._mojo.mark_dirty()
+
+    def get_focal_pixel(self) -> np.ndarray:
+        return self.mjcf.focalpixel
+
     def set_fovy(self, fovy: float):
         self.mjcf.fovy = fovy
         self._mojo.mark_dirty()

tests/elements/test_camera.py

@@ -24,6 +24,12 @@ def test_get_set_position(mojo: Mojo, camera: Camera):
     assert_array_equal(camera.get_position(), expected)
 
 
+def test_get_set_quaternion(mojo: Mojo, camera: Camera):
+    expected = np.array([0, 1, 0, 0])
+    camera.set_quaternion(expected)
+    assert_array_equal(camera.get_quaternion(), expected)
+
+
 def test_get_set_fovy(mojo: Mojo, camera: Camera):
     expected_fovy = 50.0
     camera.set_fovy(expected_fovy)
@@ -43,3 +49,10 @@ def test_get_set_sensor_size(mojo: Mojo, camera: Camera):
     camera.set_sensor_size(expected_sensor_size)
     mojo.step()  # This will recompile physics to check xml is correctly parsed
     assert_array_equal(camera.get_sensor_size(), expected_sensor_size)
+
+
+def test_get_set_focal_pixel(mojo: Mojo, camera: Camera):
+    expected_focal_pixel = np.array([400, 400], dtype=np.float64)
+    camera.set_focal_pixel(expected_focal_pixel)
+    mojo.step()  # This will recompile physics to check xml is correctly parsed
+    assert_array_equal(camera.get_focal_pixel(), expected_focal_pixel)