RSDK-4440: Add extra struct in method signatures for camera (#413)

Co-authored-by: Sean Yu <[email protected]>
Co-authored-by: Sean Yu <[email protected]>

Author: 3 people

examples/server/v1/components.py

@@ -351,13 +351,13 @@ def __init__(self, name: str):
     def __del__(self):
         self.image.close()
 
-    async def get_image(self, mime_type: str = "", **kwargs) -> Image.Image:
+    async def get_image(self, mime_type: str = "", extra: Optional[Dict[str, Any]] = None, **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]:
+    async def get_point_cloud(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Tuple[bytes, str]:
         raise NotImplementedError()
 
     async def get_properties(self, **kwargs) -> Camera.Properties:

src/viam/components/camera/camera.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Final, List, NamedTuple, Optional, Tuple, Union
+from typing import Final, List, NamedTuple, Optional, Tuple, Union, Any, Dict
 
 from PIL.Image import Image
 
@@ -35,7 +35,14 @@ class Properties(NamedTuple):
         """The distortion parameters of the camera"""
 
     @abc.abstractmethod
-    async def get_image(self, mime_type: str = "", *, timeout: Optional[float] = None, **kwargs) -> Union[Image, RawImage]:
+    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.
         Be sure to close the image when finished.
 
@@ -66,7 +73,12 @@ async def get_images(self, *, timeout: Optional[float] = None, **kwargs) -> Tupl
         ...
 
     @abc.abstractmethod
-    async def get_point_cloud(self, *, timeout: Optional[float] = None, **kwargs) -> Tuple[bytes, str]:
+    async def get_point_cloud(self,
+                              *,
+                              extra: Optional[Dict[str, Any]] = None,
+                              timeout: Optional[float] = None,
+                              **kwargs
+                              ) -> Tuple[bytes, str]:
         """
         Get the next point cloud from the camera. This will be
         returned as bytes with a mimetype describing

src/viam/components/camera/client.py

@@ -43,8 +43,15 @@ def __init__(self, name: str, channel: Channel):
         self.client = CameraServiceStub(channel)
         super().__init__(name)
 
-    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)
+    async def get_image(self,
+                        mime_type: str = "",
+                        *,
+                        extra: Optional[Dict[str, Any]] = None,
+                        timeout: Optional[float] = None
+                        ) -> Union[Image.Image, RawImage]:
+        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)
 
@@ -59,8 +66,10 @@ async def get_images(self, *, timeout: Optional[float] = None) -> Tuple[List[Nam
         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)
+    async def get_point_cloud(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> Tuple[bytes, str]:
+        if extra is None:
+            extra = {}
+        request = GetPointCloudRequest(name=self.name, mime_type=CameraMimeType.PCD, extra=dict_to_struct(extra))
         response: GetPointCloudResponse = await self.client.GetPointCloud(request, timeout=timeout)
         return (response.point_cloud, response.mime_type)
 

src/viam/components/camera/service.py

@@ -39,7 +39,7 @@ async def GetImage(self, stream: Stream[GetImageRequest, GetImageResponse]) -> N
         camera = self.get_resource(name)
 
         timeout = stream.deadline.time_remaining() if stream.deadline else None
-        image = await camera.get_image(request.mime_type, timeout=timeout, metadata=stream.metadata)
+        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:
@@ -106,7 +106,7 @@ async def GetPointCloud(self, stream: Stream[GetPointCloudRequest, GetPointCloud
         name = request.name
         camera = self.get_resource(name)
         timeout = stream.deadline.time_remaining() if stream.deadline else None
-        pc, mimetype = await camera.get_point_cloud(timeout=timeout, metadata=stream.metadata)
+        pc, mimetype = await camera.get_point_cloud(timeout=timeout, extra=struct_to_dict(request.extra), metadata=stream.metadata)
         response = GetPointCloudResponse(mime_type=mimetype, point_cloud=pc)
         await stream.send_message(response)
 

src/viam/gen/component/camera/v1/camera_pb2.pyi

@@ -359,4 +359,4 @@ class DistortionParameters(google.protobuf.message.Message):
 
     def ClearField(self, field_name: typing_extensions.Literal['model', b'model', 'parameters', b'parameters']) -> None:
         ...
-global___DistortionParameters = DistortionParameters
+global___DistortionParameters = DistortionParameters

tests/mocks/components.py

@@ -394,6 +394,7 @@ def __init__(self, name: str):
         self.image = Image.new("RGBA", (100, 100), "#AABBCCDD")
         self.geometries = GEOMETRIES
         self.point_cloud = b"THIS IS A POINT CLOUD"
+        self.extra = None
         self.props = Camera.Properties(
             False,
             IntrinsicParameters(width_px=1, height_px=2, focal_x_px=3, focal_y_px=4, center_x_px=5, center_y_px=6),
@@ -405,7 +406,13 @@ def __init__(self, name: str):
         self.metadata = ResponseMetadata(captured_at=ts)
         super().__init__(name)
 
-    async def get_image(self, mime_type: str = "", timeout: Optional[float] = None, **kwargs) -> Union[Image.Image, RawImage]:
+    async def get_image(self,
+                        mime_type: str = "",
+                        extra: Optional[Dict[str, Any]] = None,
+                        timeout: Optional[float] = None,
+                        **kwargs
+                        ) -> Union[Image.Image, RawImage]:
+        self.extra = extra
         self.timeout = timeout
         mime_type, is_lazy = CameraMimeType.from_lazy(mime_type)
         if is_lazy or (not CameraMimeType.is_supported(mime_type)):
@@ -425,7 +432,13 @@ async def get_images(self, timeout: Optional[float] = None, **kwargs) -> Tuple[L
             )
         ], self.metadata
 
-    async def get_point_cloud(self, *, timeout: Optional[float] = None, **kwargs) -> Tuple[bytes, str]:
+    async def get_point_cloud(self,
+                              *,
+                              extra: Optional[Dict[str, Any]] = None,
+                              timeout: Optional[float] = None,
+                              **kwargs
+                              ) -> Tuple[bytes, str]:
+        self.extra = extra
         self.timeout = timeout
         return self.point_cloud, CameraMimeType.PCD
 

tests/test_camera.py

@@ -85,10 +85,13 @@ def generic_service(camera: Camera) -> GenericRPCService:
 
 class TestCamera:
     @pytest.mark.asyncio
-    async def test_get_image(self, camera: Camera, image: Image.Image):
+    async def test_get_image(self, camera: MockCamera, image: Image.Image):
         img = await camera.get_image(CameraMimeType.PNG)
         assert img == image
 
+        img = await camera.get_image(CameraMimeType.PNG, {"1": 1})
+        assert camera.extra == {"1": 1}
+
         img = await camera.get_image(CameraMimeType.VIAM_RGBA)
         assert isinstance(img, Image.Image)
 
@@ -104,10 +107,13 @@ async def test_get_images(self, camera: Camera, image: Image.Image, metadata: Re
         assert md == metadata
 
     @pytest.mark.asyncio
-    async def test_get_point_cloud(self, camera: Camera, point_cloud: bytes):
+    async def test_get_point_cloud(self, camera: MockCamera, point_cloud: bytes):
         pc, _ = await camera.get_point_cloud()
         assert pc == point_cloud
 
+        await camera.get_point_cloud(extra={"1": 1})
+        assert camera.extra == {"1": 1}
+
     @pytest.mark.asyncio
     async def test_get_properties(self, camera: Camera, properties: Camera.Properties):
         props = await camera.get_properties()