RSDK-3050 - Add get_images (#358)

Co-authored-by: viambot <[email protected]>
Co-authored-by: Naveed Jooma <[email protected]>

Author: 3 people

examples/server/v1/components.py

@@ -34,6 +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.operations import run_with_operation
 from viam.proto.common import (
     AnalogStatus,
@@ -45,6 +46,7 @@
     Pose,
     PoseInFrame,
     Vector3,
+    ResponseMetadata,
 )
 from viam.proto.component.arm import JointPositions
 from viam.proto.component.audioinput import AudioChunk, AudioChunkInfo, SampleFormat
@@ -332,6 +334,9 @@ def __del__(self):
     async def get_image(self, mime_type: str = "", **kwargs) -> Image.Image:
         return self.image.copy()
 
+    async def get_images(self, timeout: Optional[float] = None, **kwargs) -> Tuple[List[NamedImage], ResponseMetadata]:
+        raise NotImplementedError()
+
     async def get_point_cloud(self, **kwargs) -> Tuple[bytes, str]:
         raise NotImplementedError()
 

poetry.lock

@@ -1,8 +1,10 @@
 import abc
-from typing import Final, NamedTuple, Optional, Tuple, Union
+from typing import Final, List, NamedTuple, Optional, Tuple, Union
 
 from PIL.Image import Image
 
+from viam.media.video import NamedImage
+from viam.proto.common import ResponseMetadata
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_COMPONENT, Subtype
 
 from ..component_base import ComponentBase
@@ -48,6 +50,21 @@ async def get_image(self, mime_type: str = "", *, timeout: Optional[float] = Non
         """
         ...
 
+    @abc.abstractmethod
+    async def get_images(self, *, timeout: Optional[float] = None, **kwargs) -> Tuple[List[NamedImage], ResponseMetadata]:
+        """Get simultaneous images from different sensors, along with associated metadata.
+        This should not be used for getting a time series of images from the same sensor.
+
+        Returns:
+            Tuple[List[NamedImage], ResponseMetadata]:
+                - List[NamedImage]:
+                  The list of images returned from the camera system.
+
+                - ResponseMetadata:
+                  The metadata associated with this response
+        """
+        ...
+
     @abc.abstractmethod
     async def get_point_cloud(self, *, timeout: Optional[float] = None, **kwargs) -> Tuple[bytes, str]:
         """

src/viam/components/camera/camera.py

@@ -1,15 +1,17 @@
 from io import BytesIO
-from typing import Mapping, Optional, Tuple, Union
+from typing import List, Mapping, Optional, Tuple, Union
 
 from grpclib.client import Channel
 from PIL import Image
 
-from viam.media.video import LIBRARY_SUPPORTED_FORMATS, CameraMimeType
-from viam.proto.common import DoCommandRequest, DoCommandResponse
+from viam.media.video import LIBRARY_SUPPORTED_FORMATS, CameraMimeType, NamedImage
+from viam.proto.common import DoCommandRequest, DoCommandResponse, ResponseMetadata
 from viam.proto.component.camera import (
     CameraServiceStub,
     GetImageRequest,
     GetImageResponse,
+    GetImagesRequest,
+    GetImagesResponse,
     GetPointCloudRequest,
     GetPointCloudResponse,
     GetPropertiesRequest,
@@ -21,6 +23,16 @@
 from . import Camera, RawImage
 
 
+def get_image_from_response(data: bytes, response_mime_type: str, request_mime_type: Optional[str] = None) -> Union[Image.Image, RawImage]:
+    if request_mime_type is None:
+        request_mime_type = response_mime_type
+    _, is_lazy = CameraMimeType.from_lazy(request_mime_type)
+    if is_lazy or not (CameraMimeType.is_supported(response_mime_type)):
+        image = RawImage(data=data, mime_type=response_mime_type)
+        return image
+    return Image.open(BytesIO(data), formats=LIBRARY_SUPPORTED_FORMATS)
+
+
 class CameraClient(Camera, ReconfigurableResourceRPCClientBase):
     """
     gRPC client for the Camera component
@@ -34,11 +46,18 @@ def __init__(self, name: str, channel: Channel):
     async def get_image(self, mime_type: str = "", *, timeout: Optional[float] = None) -> Union[Image.Image, RawImage]:
         request = GetImageRequest(name=self.name, mime_type=mime_type)
         response: GetImageResponse = await self.client.GetImage(request, timeout=timeout)
-        _, is_lazy = CameraMimeType.from_lazy(request.mime_type)
-        if is_lazy or not (CameraMimeType.is_supported(response.mime_type)):
-            image = RawImage(response.image, response.mime_type)
-            return image
-        return Image.open(BytesIO(response.image), formats=LIBRARY_SUPPORTED_FORMATS)
+        return get_image_from_response(response.image, response_mime_type=response.mime_type, request_mime_type=request.mime_type)
+
+    async def get_images(self, *, timeout: Optional[float] = None) -> Tuple[List[NamedImage], ResponseMetadata]:
+        request = GetImagesRequest(name=self.name)
+        response: GetImagesResponse = await self.client.GetImages(request, timeout=timeout)
+        imgs = []
+        for img_data in response.images:
+            mime_type = CameraMimeType.from_proto(img_data.format)
+            img = NamedImage(img_data.source_name, img_data.image, mime_type)
+            imgs.append(img)
+        resp_metadata: ResponseMetadata = response.response_metadata
+        return imgs, resp_metadata
 
     async def get_point_cloud(self, *, timeout: Optional[float] = None) -> Tuple[bytes, str]:
         request = GetPointCloudRequest(name=self.name, mime_type=CameraMimeType.PCD)

src/viam/components/camera/client.py

@@ -10,10 +10,13 @@
     CameraServiceBase,
     GetImageRequest,
     GetImageResponse,
+    GetImagesRequest,
+    GetImagesResponse,
     GetPointCloudRequest,
     GetPointCloudResponse,
     GetPropertiesRequest,
     GetPropertiesResponse,
+    Image,
     RenderFrameRequest,
 )
 from viam.resource.rpc_service_base import ResourceRPCServiceBase
@@ -61,6 +64,25 @@ async def GetImage(self, stream: Stream[GetImageRequest, GetImageResponse]) -> N
         response.image = img_bytes
         await stream.send_message(response)
 
+    async def GetImages(self, stream: Stream[GetImagesRequest, GetImagesResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        name = request.name
+        camera = self.get_resource(name)
+
+        timeout = stream.deadline.time_remaining() if stream.deadline else None
+        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()
+        response = GetImagesResponse(images=img_bytes_lst, response_metadata=metadata)
+        await stream.send_message(response)
+
     async def RenderFrame(self, stream: Stream[RenderFrameRequest, HttpBody]) -> None:
         request = await stream.recv_message()
         assert request is not None

src/viam/components/camera/service.py

@@ -499,6 +499,50 @@ class CreateOrganizationInviteResponse(google.protobuf.message.Message):
         ...
 global___CreateOrganizationInviteResponse = CreateOrganizationInviteResponse
 
+@typing_extensions.final
+class UpdateOrganizationInviteAuthorizationsRequest(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    ORGANIZATION_ID_FIELD_NUMBER: builtins.int
+    EMAIL_FIELD_NUMBER: builtins.int
+    ADD_AUTHORIZATIONS_FIELD_NUMBER: builtins.int
+    REMOVE_AUTHORIZATIONS_FIELD_NUMBER: builtins.int
+    organization_id: builtins.str
+    email: builtins.str
+
+    @property
+    def add_authorizations(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Authorization]:
+        ...
+
+    @property
+    def remove_authorizations(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Authorization]:
+        ...
+
+    def __init__(self, *, organization_id: builtins.str=..., email: builtins.str=..., add_authorizations: collections.abc.Iterable[global___Authorization] | None=..., remove_authorizations: collections.abc.Iterable[global___Authorization] | None=...) -> None:
+        ...
+
+    def ClearField(self, field_name: typing_extensions.Literal['add_authorizations', b'add_authorizations', 'email', b'email', 'organization_id', b'organization_id', 'remove_authorizations', b'remove_authorizations']) -> None:
+        ...
+global___UpdateOrganizationInviteAuthorizationsRequest = UpdateOrganizationInviteAuthorizationsRequest
+
+@typing_extensions.final
+class UpdateOrganizationInviteAuthorizationsResponse(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    INVITE_FIELD_NUMBER: builtins.int
+
+    @property
+    def invite(self) -> global___OrganizationInvite:
+        ...
+
+    def __init__(self, *, invite: global___OrganizationInvite | None=...) -> None:
+        ...
+
+    def HasField(self, field_name: typing_extensions.Literal['invite', b'invite']) -> builtins.bool:
+        ...
+
+    def ClearField(self, field_name: typing_extensions.Literal['invite', b'invite']) -> None:
+        ...
+global___UpdateOrganizationInviteAuthorizationsResponse = UpdateOrganizationInviteAuthorizationsResponse
+
 @typing_extensions.final
 class DeleteOrganizationInviteRequest(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor