fix BGR -> RGB issue

bdaiinstitute · Jul 17, 2023 · 7d9c5a6 · 7d9c5a6
1 parent 1f6e2b2
commit 7d9c5a6
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 11 deletions.
diff --git a/predicators/spot_utils/perception_utils.py b/predicators/spot_utils/perception_utils.py
@@ -235,7 +235,9 @@ def query_detic_sam(image_in: np.ndarray, classes: List[str],
 
 
 def process_image_response(
-        image_response: bosdyn.api.image_pb2.ImageResponse) -> np.ndarray:
+        image_response: bosdyn.api.image_pb2.ImageResponse,
+        to_rgb: bool = False
+) -> np.ndarray:
     """Given a Boston Dynamics SDK image response, extract the correct np array
     corresponding to the image."""
     # pylint: disable=no-member
@@ -270,6 +272,12 @@ def process_image_response(
     else:
         img = cv2.imdecode(img, -1)
 
+    # TODO also add here; to check duplicate
+    # Convert to RGB color, as some perception models assume RGB format
+    # By default, still use BGR to keep backward compability
+    if to_rgb:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
     return img
 
 

diff --git a/predicators/spot_utils/spot_utils.py b/predicators/spot_utils/spot_utils.py
@@ -284,9 +284,12 @@ def get_single_camera_image(self,
         img_req = build_image_request(
             source_name,
             quality_percent=100,
-            pixel_format=None if
-            ('hand' in source_name or 'depth' in source_name) else
-            image_pb2.Image.PIXEL_FORMAT_RGB_U8)
+            pixel_format=(
+                None if
+                ('hand' in source_name or 'depth' in source_name) else
+                image_pb2.Image.PIXEL_FORMAT_RGB_U8
+            )
+        )
         image_response = self.image_client.get_image([img_req])
 
         # Format image before detecting apriltags.
@@ -491,9 +494,14 @@ def get_sam_object_loc_from_camera(
         _, rgb_img_response = self.get_single_camera_image(source_rgb, True)
         _, depth_img_response = self.get_single_camera_image(
             source_depth, True)
+        # TODO fix - to rgb in processing response
         image = {
-            'rgb': process_image_response(rgb_img_response[0]),
-            'depth': process_image_response(depth_img_response[0]),
+            'rgb': process_image_response(
+                rgb_img_response[0], to_rgb=True
+            ),
+            'depth': process_image_response(
+                depth_img_response[0], to_rgb=True
+            ),
         }
         image_responses = {
             'rgb': rgb_img_response[0],
@@ -959,11 +967,11 @@ def arm_object_grasp(self, obj: Object) -> None:
             # stand-in for the cube, which is quite a hack.
             # We will remove this and do correct object classing
             # in a future PR
-            results = get_pixel_locations_with_detic_sam(classes=\
-                                    [obj_name_to_vision_prompt['brush']],
-                                    in_res_image=image_for_sam,
-                                    plot=\
-                                    CFG.spot_visualize_vision_model_outputs)
+            results = get_pixel_locations_with_detic_sam(
+                classes=[obj_name_to_vision_prompt['brush']],
+                in_res_image=image_for_sam,
+                plot=CFG.spot_visualize_vision_model_outputs
+            )
 
             if len(results) > 0:
                 # We only want the most likely sample (for now).