Make depth image fast

bd_spot_wrapper/spot_wrapper/redis_cache_getter.py

@@ -0,0 +1,77 @@
+import os
+
+import redis
+from bosdyn.api import image_pb2
+
+REDIS_PORT = os.environ.get("REDIS_PORT", 6379)
+hand_rgbd_key_name_in_redis = os.environ.get("hand_rgbd_key_name_for_redis", "HandRGBD")
+
+
+class RedisClient:
+    def __init__(self, host="localhost", port=None, db=0):
+        port = port or REDIS_PORT
+        self.redis_client = redis.StrictRedis(host=host, port=port, db=db)
+        print("Redis connected")
+
+    def get_latest_hand_rgbd(self):
+        """Retrieve the latest data from Redis and check if it is recent."""
+        # Get the serialized data from Redis
+        return False
+        latest_data = self.redis_client.get(f"{hand_rgbd_key_name_in_redis}")
+        if not latest_data:
+            return False  # No data available in Redis
+        image_responses_serialized = latest_data.split(b"_delimeter_")
+        image_responses = []
+        for image_response_serialized in image_responses_serialized:
+            image_response = image_pb2.ImageResponse()
+            image_response.ParseFromString(image_response_serialized)
+            image_responses.append(image_response)
+        return image_responses
+
+
+if __name__ == "__main__":
+    import time
+
+    def image_response_to_cv2(image_response, reorient=True):
+        import cv2
+        import numpy as np
+
+        if (
+            image_response.shot.image.pixel_format
+            == image_pb2.Image.PIXEL_FORMAT_DEPTH_U16
+            and image_response.shot.image.format == image_pb2.Image.FORMAT_RAW
+        ):
+            dtype = np.uint16
+        else:
+            dtype = np.uint8
+
+        # img = np.fromstring(image_response.shot.image.data, dtype=dtype)
+        img = np.frombuffer(image_response.shot.image.data, dtype=dtype)
+        if image_response.shot.image.format == image_pb2.Image.FORMAT_RAW:
+            img = img.reshape(
+                image_response.shot.image.rows, image_response.shot.image.cols
+            )
+        else:
+            img = cv2.imdecode(img, -1)
+
+        return img
+
+    redisclientforspot = RedisClient()
+
+    prev_frame_time, new_frame_time = 0, 0
+
+    while True:
+        image_responses = redisclientforspot.get_latest_hand_rgbd()
+        if image_responses:
+            # breakpoint()
+            images = [
+                image_response_to_cv2(image_response)
+                for image_response in image_responses
+            ]
+        else:
+            raise Exception("data doesn't exist yet or expired")
+        new_frame_time = time.time()  # type: ignore
+        fps = 1 / (new_frame_time - prev_frame_time)
+        prev_frame_time = new_frame_time
+        print(f"FPS {fps}")
+        time.sleep(0.1)

bd_spot_wrapper/spot_wrapper/spot.py

@@ -85,6 +85,10 @@
 from spot_rl.utils.gripper_t_intel_path import GRIPPER_T_INTEL_PATH
 from spot_rl.utils.pixel_to_3d_conversion_utils import project_3d_to_pixel_uv
 from spot_rl.utils.utils import ros_frames as rf
+try:
+    from spot_wrapper.redis_cache_getter import RedisClient
+except:
+    RedisClient = False
 from spot_wrapper.utils import (
     get_angle_between_forward_and_target,
     get_position_and_vel_values,
@@ -210,6 +214,7 @@ def __init__(self, client_name_prefix):
             RobotCommandClient.default_service_name
         )
         self.image_client = robot.ensure_client(ImageClient.default_service_name)
+        self.redis_client = RedisClient() if RedisClient else False
 
         # Make our intel image client
         try:
@@ -761,7 +766,8 @@ def grasp_point_in_image(
             camera_model=image_response.source.pinhole,
             walk_gaze_mode=3,
         )
-
+        print(f"check top_down_grasp: {top_down_grasp}")
+        print(f"check horizontal_grasp: {horizontal_grasp}")
         if top_down_grasp or horizontal_grasp:
             if top_down_grasp:
                 # Add a constraint that requests that the x-axis of the gripper is
@@ -1216,7 +1222,7 @@ def write_home_robot(self):
 
     def read_home_robot(self):
         """Returns - Tuple of global_T_home & robot_recenter_yaw. Both will be None if Home.txt file doesn't exist"""
-        print("Reading robot pose w.r.t home from home.txt")
+        # print("Reading robot pose w.r.t home from home.txt")
         global_T_home = None
         robot_recenter_yaw = None
         if osp.isfile(HOME_TXT):
@@ -1379,6 +1385,9 @@ def get_hand_image(self, is_rgb=True, force_get_gripper=False):
         ):
             return self.select_hand_image(img_src=realsense_img_srcs)
         else:
+            cache_rgbd = self.redis_client.get_latest_hand_rgbd() if self.redis_client else False
+            if cache_rgbd:
+                return cache_rgbd
             return self.select_hand_image(is_rgb=is_rgb)
 
     def get_camera_intrinsics(