From 82015ab178cb1171a5c44d655ddc2cfee64d2e57 Mon Sep 17 00:00:00 2001
From: Filippo Luca Ferretti <102977828+flferretti@users.noreply.github.com>
Date: Wed, 3 Apr 2024 15:28:05 +0200
Subject: [PATCH] Support user-defined cameras in `mujoco.loaders`

---
 src/jaxsim/mujoco/loaders.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/jaxsim/mujoco/loaders.py b/src/jaxsim/mujoco/loaders.py
index 4a3012272..bb85eb4f3 100644
--- a/src/jaxsim/mujoco/loaders.py
+++ b/src/jaxsim/mujoco/loaders.py
@@ -159,6 +159,9 @@ def convert(
         considered_joints: list[str] | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
+        cameras: (
+            list[dict[str, str, str, str, str]] | dict[str, str, str, str, str]
+        ) = None,
     ) -> tuple[str, dict[str, Any]]:
         """
         Converts a ROD model to a Mujoco MJCF string.
@@ -166,6 +169,9 @@ def convert(
         Args:
             rod_model: The ROD model to convert.
             considered_joints: The list of joint names to consider in the conversion.
+            plane_normal: The normal vector of the plane.
+            heightmap: Whether to generate a heightmap.
+            cameras: The list of cameras to add to the scene.
 
         Returns:
             tuple: A tuple containing the MJCF string and the assets dictionary.
@@ -470,6 +476,11 @@ def convert(
             fovy="60",
         )
 
+        # Add user-defined camera
+        cameras = cameras if cameras is not None else {}
+        for camera in cameras if isinstance(cameras, list) else [cameras]:
+            _ = ET.SubElement(worldbody_element, "camera", **camera)
+
         # ------------------------------------------------
         # Add a light following the  CoM of the first link
         # ------------------------------------------------
@@ -504,6 +515,9 @@ def convert(
         model_name: str | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
+        cameras: (
+            list[dict[str, str, str, str, str]] | dict[str, str, str, str, str]
+        ) = None,
     ) -> tuple[str, dict[str, Any]]:
         """
         Converts a URDF file to a Mujoco MJCF string.
@@ -512,6 +526,9 @@ def convert(
             urdf: The URDF file to convert.
             considered_joints: The list of joint names to consider in the conversion.
             model_name: The name of the model to convert.
+            plane_normal: The normal vector of the plane.
+            heightmap: Whether to generate a heightmap.
+            cameras: The list of cameras to add to the scene.
 
         Returns:
             tuple: A tuple containing the MJCF string and the assets dictionary.
@@ -530,6 +547,7 @@ def convert(
             considered_joints=considered_joints,
             plane_normal=plane_normal,
             heightmap=heightmap,
+            cameras=cameras,
         )
 
 
@@ -541,6 +559,9 @@ def convert(
         model_name: str | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
+        cameras: (
+            list[dict[str, str, str, str, str]] | dict[str, str, str, str, str]
+        ) = None,
     ) -> tuple[str, dict[str, Any]]:
         """
         Converts a SDF file to a Mujoco MJCF string.
@@ -549,6 +570,9 @@ def convert(
             sdf: The SDF file to convert.
             considered_joints: The list of joint names to consider in the conversion.
             model_name: The name of the model to convert.
+            plane_normal: The normal vector of the plane.
+            heightmap: Whether to generate a heightmap.
+            cameras: The list of cameras to add to the scene.
 
         Returns:
             tuple: A tuple containing the MJCF string and the assets dictionary.
@@ -567,4 +591,5 @@ def convert(
             considered_joints=considered_joints,
             plane_normal=plane_normal,
             heightmap=heightmap,
+            cameras=cameras,
         )