RSDK-1623 - decode custom depth mime type

src/viam/components/camera/camera.py

@@ -37,6 +37,9 @@ async def get_image(self, mime_type: str = "", *, timeout: Optional[float] = Non
         """Get the next image from the camera as an Image or RawImage.
         Be sure to close the image when finished.
 
+        NOTE: If the mime type is ``image/vnd.viam.dep`` you can use :func:`viam.media.video.RawImage.bytes_to_depth_array`
+        to convert the data to a standard representation.
+
         Args:
             mime_type (str): The desired mime type of the image. This does not guarantee output type
 

src/viam/media/video.py

@@ -1,10 +1,12 @@
+from array import array
 from enum import Enum
 from io import BytesIO
-from typing import NamedTuple, Tuple, Union
-
 from PIL.Image import Image
+from typing import List, NamedTuple, Tuple, Union
 from typing_extensions import Self
 
+from viam.errors import NotSupportedError
+
 from .viam_rgba_plugin import RGBA_FORMAT_LABEL
 
 LAZY_SUFFIX = "+lazy"
@@ -34,6 +36,27 @@ def close(self):
         """Close the image and release resources. For RawImage, this is a noop."""
         return
 
+    def bytes_to_depth_array(self) -> List[List[int]]:
+        """Decode the data of an image that has the custom depth MIME type ``image/vnd.viam.dep`` into
+        a standard representation.
+
+        Raises:
+            NotSupportedError: Raised if given an image that is not of MIME type `image/vnd.viam.dep`.
+
+        Returns:
+            List[List[int]]: The standard representation of the image.
+        """
+        if self.mime_type != "image/vnd.viam.dep":
+            raise NotSupportedError("Type must be `image/vnd.viam.dep` to use bytes_to_depth_array()")
+
+        width = int.from_bytes(self.data[8:16], "big")
+        height = int.from_bytes(self.data[16:24], "big")
+        depth_arr = array("H", self.data[24:])
+        depth_arr.byteswap()
+
+        depth_arr_2d = [[depth_arr[row * width + col] for col in range(width)] for row in range(height)]
+        return depth_arr_2d
+
 
 class CameraMimeType(str, Enum):
     VIAM_RGBA = "image/vnd.viam.rgba"

tests/test_media.py

@@ -0,0 +1,23 @@
+from array import array
+import os
+import pytest
+
+from viam.media.video import RawImage
+
+
+class TestRawImage:
+    @pytest.mark.asyncio
+    async def test_bytes_to_depth_array(self):
+        with open(f"{os.path.dirname(__file__)}/../data/fakeDM.vnd.viam.dep", "rb") as depth_map:
+            image = RawImage(depth_map.read(), "image/vnd.viam.dep")
+        assert isinstance(image, RawImage)
+        standard_data = image.bytes_to_depth_array()
+        assert len(standard_data) == 10
+        assert len(standard_data[0]) == 20
+        data_arr = array("H", image.data[24:])
+        data_arr.byteswap()
+        assert len(data_arr) == 200
+        assert standard_data[0][0] == data_arr[0]
+        assert standard_data[-1][3] == data_arr[183]
+        assert standard_data[-1][3] == 9 * 3
+        assert standard_data[4][4] == 4 * 4