Add fix to also extract labels with python Bagtransformer

leggedrobotics · Aug 23, 2023 · 018310b · 018310b
1 parent 42055b6
commit 018310b
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/wild_visual_navigation_ros/src/wild_visual_navigation_ros/ros_converter.py b/wild_visual_navigation_ros/src/wild_visual_navigation_ros/ros_converter.py
@@ -106,7 +106,16 @@ def ros_tf_to_torch(tf_pose, device="cpu"):
 
 
 def ros_image_to_torch(ros_img, desired_encoding="rgb8", device="cpu"):
-    if isinstance(ros_img, Image):
+    if type(ros_img).__name__ is "_sensor_msgs__Image":
+        np_image = CV_BRIDGE.imgmsg_to_cv2(ros_img, desired_encoding=desired_encoding)
+
+    elif type(ros_img).__name__ is "_sensor_msgs__CompressedImage":
+        np_arr = np.fromstring(ros_img.data, np.uint8)
+        np_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+        if "bgr" in ros_img.format:
+            np_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
+
+    elif isinstance(ros_img, Image):
         np_image = CV_BRIDGE.imgmsg_to_cv2(ros_img, desired_encoding=desired_encoding)
 
     elif isinstance(ros_img, CompressedImage):