Merge pull request #14 from agimus-project/cosypose_model_type

add choice between model_type pbr and synth+real
agimus-project · Oct 21, 2024 · 2e76d43 · 2e76d43
2 parents cdfd216 + 8e41f9a
commit 2e76d43
Show file tree

Hide file tree

Showing 10 changed files with 45 additions and 20 deletions.
diff --git a/happypose_examples/config/cosypose_params.yaml b/happypose_examples/config/cosypose_params.yaml
@@ -9,16 +9,16 @@
     # Specifies which pose estimator to use in the pipeline
     pose_estimator_type: "cosypose"
     cosypose:
-      # Name of a dataset used during training
+      # Name of BOP dataset, used to load specific weights and object models
       dataset_name: $(var dataset_name)
+      # Type of neural network model to use. Available 'pbr'|'synth+real'
+      model_type: $(var model_type)
       # Object renderer parameters
       renderer:
         # Specifies which renderer to use in the pipeline
         renderer_type: "panda3d"
         # Number of CPU cores to use during rendering
         n_workers: 8
-        # Render objects on a GPU
-        gpu_renderer: false
       # Parameters expected on the runtime
       inference:
         # 2D image detection parameters

diff --git a/happypose_examples/config/cosypose_params_multiview.yaml b/happypose_examples/config/cosypose_params_multiview.yaml
@@ -15,16 +15,16 @@
     # Specifies which pose estimator to use in the pipeline
     pose_estimator_type: "cosypose"
     cosypose:
-      # Name of a dataset used during training
+      # Name of BOP dataset, used to load specific weights and object models
       dataset_name: $(var dataset_name)
+      # Type of neural network model to use. Available 'pbr'|'synth+real'
+      model_type: $(var model_type)
       # Object renderer parameters
       renderer:
         # Specifies which renderer to use in the pipeline
         renderer_type: "panda3d"
         # Number of CPU cores to use during rendering
         n_workers: 8
-        # Render objects on a GPU
-        gpu_renderer: false
       # Parameters expected on the runtime
       inference:
         # 2D image detection parameters

diff --git a/happypose_examples/launch/common.launch.py b/happypose_examples/launch/common.launch.py
@@ -68,7 +68,12 @@ def generate_launch_description():
         DeclareLaunchArgument(
             "dataset_name",
             default_value="ycbv",
-            description="Which dataset to use for inference.",
+            description="Name of BOP dataset, used to load specific weights and object models.",
+        ),
+        DeclareLaunchArgument(
+            "model_type",
+            default_value="pbr",
+            description="Type of neural network model to use. Available: 'pbr'|'synth+real'.",
         ),
         DeclareLaunchArgument(
             "device",

diff --git a/happypose_examples/launch/multi_view_demo.launch.py b/happypose_examples/launch/multi_view_demo.launch.py
@@ -82,6 +82,7 @@ def launch_setup(
         ),
         launch_arguments={
             "dataset_name": LaunchConfiguration("dataset_name"),
+            "model_type": LaunchConfiguration("model_type"),
             "device": LaunchConfiguration("device"),
             "use_rviz": LaunchConfiguration("use_rviz"),
             "rviz_config_path": rviz_config_path,
@@ -97,7 +98,12 @@ def generate_launch_description():
         DeclareLaunchArgument(
             "dataset_name",
             default_value="ycbv",
-            description="Which dataset to use for inference.",
+            description="Name of BOP dataset, used to load specific weights and object models.",
+        ),
+        DeclareLaunchArgument(
+            "model_type",
+            default_value="pbr",
+            description="What type of NN model to use, type of training data used, 'pbr'|'synth+real'.",
         ),
         DeclareLaunchArgument(
             "device",

diff --git a/happypose_examples/launch/single_view_demo.launch.py b/happypose_examples/launch/single_view_demo.launch.py
@@ -59,6 +59,7 @@ def launch_setup(
         ),
         launch_arguments={
             "dataset_name": LaunchConfiguration("dataset_name"),
+            "model_type": LaunchConfiguration("model_type"),
             "device": LaunchConfiguration("device"),
             "use_rviz": LaunchConfiguration("use_rviz"),
         }.items(),
@@ -72,7 +73,12 @@ def generate_launch_description():
         DeclareLaunchArgument(
             "dataset_name",
             default_value="ycbv",
-            description="Which dataset to use for inference.",
+            description="Name of BOP dataset, used to load specific weights and object models.",
+        ),
+        DeclareLaunchArgument(
+            "model_type",
+            default_value="pbr",
+            description="Type of neural network model to use. Available: 'pbr'|'synth+real'.",
         ),
         DeclareLaunchArgument(
             "device",

diff --git a/happypose_ros/happypose_ros/happypose_ros_parameters.yaml b/happypose_ros/happypose_ros/happypose_ros_parameters.yaml
@@ -44,10 +44,17 @@ happypose_ros:
     dataset_name:
       type: string
       default_value: ""
-      description: "Name of a dataset used during training."
+      description: "Name of BOP dataset, used to load specific weights and object models."
       validation:
         one_of<>: [["hope", "tless", "ycbv"]]
       read_only: true
+    model_type:
+      type: string
+      default_value: "pbr"
+      description: "Type of neural network model to use. Available: 'pbr'|'synth+real'"
+      validation:
+        one_of<>: [["pbr", "synth+real"]]
+      read_only: true
     renderer:
       renderer_type:
         type: string
@@ -63,11 +70,6 @@ happypose_ros:
         validation:
           gt_eq<>: 1
         read_only: true
-      gpu_renderer:
-        type: bool
-        default_value: true
-        description: "Render objects on a GPU."
-        read_only: true
     inference:
       detector:
         detection_th:

diff --git a/happypose_ros/happypose_ros/inference_pipeline.py b/happypose_ros/happypose_ros/inference_pipeline.py
@@ -39,6 +39,7 @@ def __init__(self, params: dict) -> None:
         # Currently only cosypose is supported
         self._wrapper = CosyPoseWrapper(
             dataset_name=self._params["cosypose"]["dataset_name"],
+            model_type=self._params["cosypose"]["model_type"],
             **self._params["cosypose"]["renderer"],
         )
 

diff --git a/happypose_ros/test/test_initialization_params.py b/happypose_ros/test/test_initialization_params.py
@@ -38,14 +38,15 @@ def happypose_params(request: pytest.FixtureRequest) -> dict:
 
 
 @pytest.fixture()
-def minial_overwrites() -> List[Parameter]:
+def minimal_overwrites() -> List[Parameter]:
     """Crete set minimal set of ROS parameters needed for the ROS node to start correctly.
 
     :return: Default, minimal parameters for the ROS node to start.
     :rtype: List[rclpy.Parameter]
     """
     return [
         Parameter("cosypose.dataset_name", Parameter.Type.STRING, "ycbv"),
+        Parameter("cosypose.model_type", Parameter.Type.STRING, "pbr"),
         Parameter("camera_names", Parameter.Type.STRING_ARRAY, ["cam_1"]),
         Parameter("cameras.cam_1.leading", Parameter.Type.BOOL, True),
         Parameter("cameras.cam_1.publish_tf", Parameter.Type.BOOL, False),
@@ -58,6 +59,7 @@ def test_no_leading(happypose_params: dict) -> None:
             **happypose_params,
             parameter_overrides=[
                 Parameter("cosypose.dataset_name", Parameter.Type.STRING, "ycbv"),
+                Parameter("cosypose.model_type", Parameter.Type.STRING, "pbr"),
                 Parameter(
                     "camera_names", Parameter.Type.STRING_ARRAY, ["cam_1", "cam_2"]
                 ),
@@ -69,11 +71,11 @@ def test_no_leading(happypose_params: dict) -> None:
 
 
 @mock.patch("multiprocessing.context.SpawnContext.Process")
-def test_minimal(happypose_params: dict, minial_overwrites: List[Parameter]) -> None:
+def test_minimal(happypose_params: dict, minimal_overwrites: List[Parameter]) -> None:
     # Check if node starts correctly with minimal number of required parameters
     happypose_node = HappyPoseNode(
         **happypose_params,
-        parameter_overrides=minial_overwrites,
+        parameter_overrides=minimal_overwrites,
     )
     happypose_node.destroy_node()
 
@@ -84,6 +86,7 @@ def test_multiple_leading(happypose_params: dict) -> None:
             **happypose_params,
             parameter_overrides=[
                 Parameter("cosypose.dataset_name", Parameter.Type.STRING, "ycbv"),
+                Parameter("cosypose.model_type", Parameter.Type.STRING, "pbr"),
                 Parameter(
                     "camera_names", Parameter.Type.STRING_ARRAY, ["cam_1", "cam_2"]
                 ),
@@ -114,13 +117,13 @@ def test_leading_publish_tf(happypose_params: dict) -> None:
 
 
 def test_device_unknown(
-    happypose_params: dict, minial_overwrites: List[Parameter]
+    happypose_params: dict, minimal_overwrites: List[Parameter]
 ) -> None:
     with pytest.raises(ParameterException) as excinfo:
         HappyPoseNode(
             **happypose_params,
             parameter_overrides=[
-                *minial_overwrites,
+                *minimal_overwrites,
                 # CUDA -1 is an invalid device
                 Parameter("device", Parameter.Type.STRING, "cuda:-1"),
             ],

diff --git a/happypose_ros/test/test_multi_view.yaml b/happypose_ros/test/test_multi_view.yaml
@@ -4,6 +4,7 @@
     pose_estimator_type: "cosypose"
     cosypose:
       dataset_name: "ycbv"
+      model_type: "pbr"
       # reduce memory footprint during tests by reducing number of workers
       renderer:
         n_workers: 1

diff --git a/happypose_ros/test/test_single_view.yaml b/happypose_ros/test/test_single_view.yaml
@@ -6,6 +6,7 @@
     visualization.publish_markers: true
     cosypose:
       dataset_name: "ycbv"
+      model_type: "pbr"
       # reduce memory footprint during tests by reducing number of workers
       renderer:
         n_workers: 1