RSDK-4089: use viam image (#567)

docs/examples/example.ipynb

@@ -161,11 +161,13 @@
    "source": [
     "from viam.components.camera import Camera\n",
     "from viam.media.video import CameraMimeType\n",
+    "from viam.media.utils.pil import viam_to_pil_image\n",
     "\n",
     "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",
+    "pil = viam_to_pil_image(image)\n",
+    "pil.save(\"foo.png\")\n",
     "\n",
     "# Don't forget to close the robot when you're done!\n",
     "await robot.close()"
@@ -1489,7 +1491,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6 (main, Oct  2 2023, 20:46:14) [Clang 14.0.3 (clang-1403.0.22.14.1)]"
+   "version": "3.11.9 (main, Apr  2 2024, 08:25:04) [Clang 15.0.0 (clang-1500.1.0.2.5)]"
   },
   "vscode": {
    "interpreter": {

examples/server/v1/components.py

@@ -10,6 +10,7 @@
     from typing import AsyncIterator
 
 from datetime import timedelta
+from io import BytesIO
 from multiprocessing import Lock
 from pathlib import Path
 from typing import Any, Dict, List, Mapping, Optional, Tuple
@@ -33,7 +34,7 @@
 from viam.errors import ResourceNotFoundError
 from viam.media import MediaStreamWithIterator
 from viam.media.audio import Audio, AudioStream
-from viam.media.video import NamedImage
+from viam.media.video import CameraMimeType, NamedImage, ViamImage
 from viam.operations import run_with_operation
 from viam.proto.common import (
     AnalogStatus,
@@ -330,14 +331,17 @@ async def get_geometries(self, extra: Optional[Dict[str, Any]] = None, **kwargs)
 class ExampleCamera(Camera):
     def __init__(self, name: str):
         p = Path(__file__)
-        self.image = Image.open(p.parent.absolute().joinpath("viam.webp"))
+        img = Image.open(p.parent.absolute().joinpath("viam.webp"))
+        buf = BytesIO()
+        img.copy().save(buf, format="JPEG")
+        self.image = ViamImage(buf.getvalue(), CameraMimeType.JPEG)
         super().__init__(name)
 
     def __del__(self):
         self.image.close()
 
-    async def get_image(self, mime_type: str = "", extra: Optional[Dict[str, Any]] = None, **kwargs) -> Image.Image:
-        return self.image.copy()
+    async def get_image(self, mime_type: str = "", extra: Optional[Dict[str, Any]] = None, **kwargs) -> ViamImage:
+        return self.image
 
     async def get_images(self, timeout: Optional[float] = None, **kwargs) -> Tuple[List[NamedImage], ResponseMetadata]:
         raise NotImplementedError()

src/viam/components/camera/__init__.py

@@ -1,4 +1,4 @@
-from viam.media.video import RawImage, ViamImage
+from viam.media.video import ViamImage
 from viam.proto.component.camera import DistortionParameters, IntrinsicParameters
 from viam.resource.registry import Registry, ResourceRegistration
 
@@ -10,7 +10,6 @@
     "Camera",
     "IntrinsicParameters",
     "DistortionParameters",
-    "RawImage",
     "ViamImage",
 ]
 

src/viam/components/camera/camera.py

@@ -1,16 +1,13 @@
 import abc
 import sys
-from typing import Any, Dict, Final, List, Optional, Tuple, Union
+from typing import Any, Dict, Final, List, 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.proto.component.camera import GetPropertiesResponse
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT, Subtype
 
 from ..component_base import ComponentBase
-from . import RawImage
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
@@ -40,8 +37,8 @@ class Camera(ComponentBase):
     @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.ViamImage.bytes_to_depth_array`
@@ -62,7 +59,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,
@@ -19,17 +17,7 @@
 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
-
-
-def get_image_from_response(data: bytes, response_mime_type: str, request_mime_type: str) -> Union[Image.Image, RawImage]:
-    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)
+from . import Camera
 
 
 class CameraClient(Camera, ReconfigurableResourceRPCClientBase):
@@ -49,12 +37,12 @@ async def get_image(
         extra: Optional[Dict[str, Any]] = None,
         timeout: Optional[float] = None,
         **__,
-    ) -> Union[Image.Image, RawImage]:
+    ) -> ViamImage:
         if extra is None:
             extra = {}
         request = GetImageRequest(name=self.name, mime_type=mime_type, extra=dict_to_struct(extra))
         response: GetImageResponse = await self.client.GetImage(request, timeout=timeout)
-        return get_image_from_response(response.image, response_mime_type=response.mime_type, request_mime_type=request.mime_type)
+        return ViamImage(response.image, CameraMimeType.from_string(response.mime_type))
 
     async def get_images(
         self,

src/viam/components/camera/service.py

@@ -1,11 +1,8 @@
 # TODO: Update type checking based with RSDK-4089
 # pyright: reportGeneralTypeIssues=false
-from typing import Dict
-
 from google.api.httpbody_pb2 import HttpBody
 from grpclib.server import Stream
 
-from viam.media.video import CameraMimeType
 from viam.proto.common import DoCommandRequest, DoCommandResponse, GetGeometriesRequest, GetGeometriesResponse
 from viam.proto.component.camera import (
     CameraServiceBase,
@@ -23,7 +20,7 @@
 from viam.resource.rpc_service_base import ResourceRPCServiceBase
 from viam.utils import dict_to_struct, struct_to_dict
 
-from . import Camera, RawImage
+from . import Camera
 
 
 class CameraRPCService(CameraServiceBase, ResourceRPCServiceBase[Camera]):
@@ -32,7 +29,6 @@ class CameraRPCService(CameraServiceBase, ResourceRPCServiceBase[Camera]):
     """
 
     RESOURCE_TYPE = Camera
-    _camera_mime_types: Dict[str, CameraMimeType] = {}
 
     async def GetImage(self, stream: Stream[GetImageRequest, GetImageResponse]) -> None:
         request = await stream.recv_message()
@@ -42,23 +38,7 @@ async def GetImage(self, stream: Stream[GetImageRequest, GetImageResponse]) -> N
 
         timeout = stream.deadline.time_remaining() if stream.deadline else None
         image = await camera.get_image(request.mime_type, extra=struct_to_dict(request.extra), timeout=timeout, metadata=stream.metadata)
-        try:
-            if not request.mime_type:
-                if camera.name not in self._camera_mime_types:
-                    self._camera_mime_types[camera.name] = CameraMimeType.JPEG
-
-                request.mime_type = self._camera_mime_types[camera.name]
-
-            mimetype, _ = CameraMimeType.from_lazy(request.mime_type)
-            if CameraMimeType.is_supported(mimetype):
-                response_mime = mimetype
-            else:
-                response_mime = request.mime_type
-            response = GetImageResponse(mime_type=response_mime)
-            img_bytes = mimetype.encode_image(image)
-        finally:
-            image.close()
-        response.image = img_bytes
+        response = GetImageResponse(mime_type=image.mime_type, image=image.data)
         await stream.send_message(response)
 
     async def GetImages(self, stream: Stream[GetImagesRequest, GetImagesResponse]) -> None:
@@ -71,12 +51,9 @@ async def GetImages(self, stream: Stream[GetImagesRequest, GetImagesResponse]) -
         images, metadata = await camera.get_images(timeout=timeout, metadata=stream.metadata)
         img_bytes_lst = []
         for img in images:
-            try:
-                fmt = img.mime_type.to_proto()
-                img_bytes = img.data
-                img_bytes_lst.append(Image(source_name=name, format=fmt, image=img_bytes))
-            finally:
-                img.close()
+            fmt = img.mime_type.to_proto()
+            img_bytes = img.data
+            img_bytes_lst.append(Image(source_name=name, format=fmt, image=img_bytes))
         response = GetImagesResponse(images=img_bytes_lst, response_metadata=metadata)
         await stream.send_message(response)
 
@@ -85,17 +62,9 @@ async def RenderFrame(self, stream: Stream[RenderFrameRequest, HttpBody]) -> Non
         assert request is not None
         name = request.name
         camera = self.get_resource(name)
-        try:
-            mimetype = CameraMimeType(request.mime_type)
-        except ValueError:
-            mimetype = CameraMimeType.JPEG
         timeout = stream.deadline.time_remaining() if stream.deadline else None
-        image = await camera.get_image(mimetype, timeout=timeout, metadata=stream.metadata)
-        try:
-            img = mimetype.encode_image(image)
-        finally:
-            image.close()
-        response = HttpBody(data=img, content_type=image.mime_type if isinstance(image, RawImage) else mimetype)  # type: ignore
+        image = await camera.get_image(request.mime_type, timeout=timeout, metadata=stream.metadata)
+        response = HttpBody(data=image.data, content_type=image.mime_type)  # type: ignore
         await stream.send_message(response)
 
     async def GetPointCloud(self, stream: Stream[GetPointCloudRequest, GetPointCloudResponse]) -> None:

src/viam/media/utils/pil.py

@@ -0,0 +1,46 @@
+from io import BytesIO
+
+from PIL import Image
+
+from viam.media.video import CameraMimeType, ViamImage
+
+from ..viam_rgba_plugin import RGBA_FORMAT_LABEL
+
+# Formats that are supported by PIL
+LIBRARY_SUPPORTED_FORMATS = ["JPEG", "PNG", RGBA_FORMAT_LABEL]
+
+
+def viam_to_pil_image(image: ViamImage) -> Image.Image:
+    """
+    Convert a ViamImage to a PIL.Image.
+
+    Args:
+        image (ViamImage): The image to convert.
+
+    Returns:
+        Image.Image: The resulting PIL.Image
+    """
+    return Image.open(BytesIO(image.data), formats=LIBRARY_SUPPORTED_FORMATS)
+
+
+def pil_to_viam_image(image: Image.Image, mime_type: CameraMimeType) -> ViamImage:
+    """
+    Convert a PIL.Image to a ViamImage.
+
+    Args:
+        image (Image.Image): The image to convert.
+        mime_type (CameraMimeType): The mime type to convert the image to.
+
+    Returns:
+        ViamImage: The resulting ViamImage
+    """
+    if mime_type.name in LIBRARY_SUPPORTED_FORMATS:
+        buf = BytesIO()
+        if image.mode == "RGBA" and mime_type == CameraMimeType.JPEG:
+            image = image.convert("RGB")
+        image.save(buf, format=mime_type.name)
+        data = buf.getvalue()
+    else:
+        raise ValueError(f"Cannot encode image to {mime_type}")
+
+    return ViamImage(data, mime_type)