Added both options

leggedrobotics · Aug 23, 2023 · 7f9b2b3 · 7f9b2b3
1 parent 989e56a
commit 7f9b2b3
Showing 1 changed file with 12 additions and 0 deletions.
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,6 +106,7 @@ def ros_tf_to_torch(tf_pose, device="cpu"):
 
 
 def ros_image_to_torch(ros_img, desired_encoding="rgb8", device="cpu"):
+    # Use for replaying bag with BagTfTransformer package
     if type(ros_img).__name__ is "_sensor_msgs__Image":
         np_image = CV_BRIDGE.imgmsg_to_cv2(ros_img, desired_encoding=desired_encoding)
 
@@ -114,6 +115,17 @@ def ros_image_to_torch(ros_img, desired_encoding="rgb8", device="cpu"):
         np_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
         if "bgr" in ros_img.format:
             np_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
+
+    # Used for replaying with rosbags via terminal
+    elif isinstance(ros_img, Image):
+        np_image = CV_BRIDGE.imgmsg_to_cv2(ros_img, desired_encoding=desired_encoding)
+
+    elif isinstance(ros_img, 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)
+
     return TO_TENSOR(np_image).to(device)