Updates on loading / saving depth images

camera_calibration/record_robot_joints.py

@@ -4,11 +4,11 @@
 import numpy as np
 import simplejson as json
 
-from gprs import config_root
-from gprs.franka_interface import FrankaInterface
-from gprs.utils import YamlConfig
-from gprs.utils.input_utils import input2action
-from gprs.utils.io_devices import SpaceMouse
+from deoxys import config_root
+from deoxys.franka_interface import FrankaInterface
+from deoxys.utils import YamlConfig
+from deoxys.utils.input_utils import input2action
+from deoxys.utils.io_devices import SpaceMouse
 
 config_folder = os.path.join(os.path.expanduser("~/"), ".deoxys_vision/calibration_configuration")
 os.makedirs(os.path.join(os.path.expanduser("~/"), config_folder), exist_ok=True)
@@ -22,7 +22,7 @@ def main():
     # print(config_root)
     robot_interface = FrankaInterface(config_root + "/charmander.yml", use_visualizer=False)
     controller_cfg = YamlConfig(config_root + "/compliant-joint-impedance-controller.yml").as_easydict()
-    control_type = "JOINT_IMPEDANCE"
+    controller_type = "JOINT_IMPEDANCE"
 
     # # Make it low impedance so that we can easily move the arm around
     # controller_cfg["Kp"]["translation"] = 50
@@ -35,7 +35,7 @@ def main():
     while True:
         spacemouse_action, grasp = input2action(
             device=device,
-            control_type="OSC_POSE",
+            controller_type="OSC_POSE",
         )
 
         if spacemouse_action is None:
@@ -49,15 +49,15 @@ def main():
                 for _ in range(5):
                     spacemouse_action, grasp = input2action(
                         device=device,
-                        control_type=control_type,
+                        controller_type=controller_type,
                     )
             elif spacemouse_action[-1] < 0:
                 recorded_joint = False
         else:
             continue
         action = list(robot_interface._state_buffer[-1].q) + [-1]
         robot_interface.control(
-            control_type=control_type, action=action, controller_cfg=controller_cfg
+            controller_type=controller_type, action=action, controller_cfg=controller_cfg
         )
 
     save_joints = []

deoxys_vision/utils/img_utils.py

@@ -57,7 +57,29 @@ def load_depth(depth_img_name):
     return cv2.imread(depth_img_name, cv2.IMREAD_UNCHANGED)
 
 def save_depth(depth_img_name, depth_img):
-    cv2.imwrite(depth_img_name, cv2.cvtColor(depth_img, cv2.CV_16U))    
+    assert(depth_img_name.endswith(".tiff")), "You are not using tiff file for saving uint16 data. Things will be screwed."
+    cv2.imwrite(depth_img_name, cv2.cvtColor(depth_img, cv2.CV_16U))
+
+def load_depth_in_rgb(depth_img_name):
+    rgb_img = cv2.imread(depth_img_name).astype(np.uint8)
+
+    depth_img = np.zeros((rgb_img.shape[0], rgb_img.shape[1])).astype(np.uint16)
+    depth_img = rgb_img[..., 1].astype(np.uint16) << 8 | rgb_img[..., 2].astype(np.uint16)
+
+    return depth_img
+
+def save_depth_in_rgb(depth_img_name, depth_img):
+    """
+    Saving depth image in the format of rgb images. The nice thing is that we can leverage the efficient PNG encoding to save almost 50% spaces compared to using tiff.
+    """
+    assert(depth_img.dtype == np.uint16)
+    assert(depth_img_name.endswith(".png")), "You are not using lossless saving. Depth image will be messed up if you want to use rgb format."
+    higher_bytes = depth_img >> 8
+    lower_bytes = depth_img & 0xFF
+    depth_rgb_img = np.zeros((depth_img.shape[0], depth_img.shape[1], 3)).astype(np.uint8)
+    depth_rgb_img[..., 1] = higher_bytes.astype(np.uint8)
+    depth_rgb_img[..., 2] = lower_bytes.astype(np.uint8)
+    cv2.imwrite(depth_img_name, depth_rgb_img)
 
 def preprocess_color(color_img, flip_channel=True):
     if flip_channel:
@@ -67,4 +89,4 @@ def preprocess_color(color_img, flip_channel=True):
 
 def preprocess_depth(depth_img):
     return np.ascontiguousarray(depth_img)
-
+