RSDK-4089: change camera api return types (#495)

Author: purplenicole730

docs/examples/example.ipynb

@@ -164,7 +164,7 @@
     "robot = await connect_with_channel()\n",
     "camera = Camera.from_robot(robot, \"camera0\")\n",
     "image = await camera.get_image(CameraMimeType.JPEG)\n",
-    "image.save(\"foo.png\")\n",
+    "image.image.save(\"foo.png\")\n",
     "\n",
     "# Don't forget to close the robot when you're done!\n",
     "await robot.close()\n"

src/viam/components/camera/camera.py

@@ -1,14 +1,12 @@
 import abc
-from typing import Any, Dict, Final, List, NamedTuple, Optional, Tuple, Union
+from typing import Any, Dict, Final, List, NamedTuple, Optional, Tuple
 
-from PIL.Image import Image
-
-from viam.media.video import NamedImage
+from viam.media.video import NamedImage, ViamImage
 from viam.proto.common import ResponseMetadata
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT, Subtype
 
 from ..component_base import ComponentBase
-from . import DistortionParameters, IntrinsicParameters, RawImage
+from . import DistortionParameters, IntrinsicParameters
 
 
 class Camera(ComponentBase):
@@ -37,8 +35,8 @@ class Properties(NamedTuple):
     @abc.abstractmethod
     async def get_image(
         self, mime_type: str = "", *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None, **kwargs
-    ) -> Union[Image, RawImage]:
-        """Get the next image from the camera as an Image or RawImage.
+    ) -> ViamImage:
+        """Get the next image from the camera as a ViamImage.
         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`
@@ -48,7 +46,7 @@ async def get_image(
             mime_type (str): The desired mime type of the image. This does not guarantee output type
 
         Returns:
-            Image | RawImage: The frame
+            ViamImage: The frame
         """
         ...
 

src/viam/components/camera/client.py

@@ -1,10 +1,8 @@
-from io import BytesIO
-from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, List, Mapping, Optional, Tuple
 
 from grpclib.client import Channel
-from PIL import Image
 
-from viam.media.video import LIBRARY_SUPPORTED_FORMATS, CameraMimeType, NamedImage
+from viam.media.video import CameraMimeType, NamedImage, ViamImage
 from viam.proto.common import DoCommandRequest, DoCommandResponse, Geometry, ResponseMetadata
 from viam.proto.component.camera import (
     CameraServiceStub,
@@ -20,17 +18,14 @@
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
 from viam.utils import ValueTypes, dict_to_struct, get_geometries, struct_to_dict
 
-from . import Camera, RawImage
+from . import Camera
 
 
-def get_image_from_response(data: bytes, response_mime_type: str, request_mime_type: str) -> Union[Image.Image, RawImage]:
+def get_image_from_response(data: bytes, response_mime_type: str, request_mime_type: str) -> ViamImage:
     if not request_mime_type:
         request_mime_type = response_mime_type
-    mime_type, is_lazy = CameraMimeType.from_lazy(request_mime_type)
-    if is_lazy or mime_type._should_be_raw:
-        image = RawImage(data=data, mime_type=response_mime_type)
-        return image
-    return Image.open(BytesIO(data), formats=LIBRARY_SUPPORTED_FORMATS)
+    mime_type, _ = CameraMimeType.from_lazy(request_mime_type)
+    return ViamImage(data, mime_type)
 
 
 class CameraClient(Camera, ReconfigurableResourceRPCClientBase):
@@ -43,9 +38,7 @@ def __init__(self, name: str, channel: Channel):
         self.client = CameraServiceStub(channel)
         super().__init__(name)
 
-    async def get_image(
-        self, mime_type: str = "", *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None
-    ) -> Union[Image.Image, RawImage]:
+    async def get_image(self, mime_type: str = "", *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> ViamImage:
         if extra is None:
             extra = {}
         request = GetImageRequest(name=self.name, mime_type=mime_type, extra=dict_to_struct(extra))

src/viam/components/camera/service.py

@@ -47,7 +47,7 @@ async def GetImage(self, stream: Stream[GetImageRequest, GetImageResponse]) -> N
 
                 request.mime_type = self._camera_mime_types[camera.name]
 
-            mimetype, is_lazy = CameraMimeType.from_lazy(request.mime_type)
+            mimetype, _ = CameraMimeType.from_lazy(request.mime_type)
             if CameraMimeType.is_supported(mimetype):
                 response_mime = mimetype
             else:

tests/test_camera.py

@@ -10,7 +10,7 @@
 from viam.components.camera import Camera, CameraClient
 from viam.components.camera.service import CameraRPCService
 from viam.components.generic.service import GenericRPCService
-from viam.media.video import LIBRARY_SUPPORTED_FORMATS, CameraMimeType, NamedImage, RawImage
+from viam.media.video import LIBRARY_SUPPORTED_FORMATS, CameraMimeType, NamedImage, RawImage, ViamImage
 from viam.proto.common import DoCommandRequest, DoCommandResponse, GetGeometriesRequest, GetGeometriesResponse, ResponseMetadata
 from viam.proto.component.camera import (
     CameraServiceStub,
@@ -253,27 +253,29 @@ async def test_get_image(self, camera: MockCamera, service: CameraRPCService, im
 
             # Test known mime type
             png_img = await client.get_image(timeout=1.82, mime_type=CameraMimeType.PNG)
-            assert isinstance(png_img, Image.Image)
-            assert png_img.tobytes() == image.tobytes()
+            assert isinstance(png_img.image, Image.Image)
+            assert png_img.image.tobytes() == image.tobytes()
             assert camera.timeout == loose_approx(1.82)
 
             # Test raw mime type
             rgba_img = await client.get_image(CameraMimeType.VIAM_RGBA)
-            assert isinstance(rgba_img, Image.Image)
-            rgba_bytes = rgba_img.tobytes()
+            assert isinstance(rgba_img.image, Image.Image)
+            rgba_bytes = rgba_img.image.tobytes()
             assert rgba_bytes == image.copy().convert("RGBA").tobytes()
 
             # Test lazy mime type
             raw_img = await client.get_image(CameraMimeType.PNG.with_lazy_suffix)
-            assert isinstance(raw_img, RawImage)
+            assert isinstance(raw_img, ViamImage)
+            assert raw_img.image is None
             assert raw_img.data == image.tobytes()
             assert raw_img.mime_type == CameraMimeType.PNG
 
             # Test unknown mime type
             raw_img = await client.get_image("unknown")
-            assert isinstance(raw_img, RawImage)
+            assert isinstance(raw_img, ViamImage)
+            assert raw_img.image is None
             assert raw_img.data == image.tobytes()
-            assert raw_img.mime_type == "unknown"
+            assert raw_img.mime_type == CameraMimeType.UNSUPPORTED
 
     @pytest.mark.asyncio
     async def test_get_images(self, camera: MockCamera, service: CameraRPCService, image: Image.Image, metadata: ResponseMetadata):