Skip to content

Commit

Permalink
RSDK-4089: change camera api return types (#495)
Browse files Browse the repository at this point in the history
  • Loading branch information
purplenicole730 authored Dec 8, 2023
1 parent fb2ee20 commit 35c64cb
Show file tree
Hide file tree
Showing 5 changed files with 25 additions and 32 deletions.
2 changes: 1 addition & 1 deletion docs/examples/example.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
14 changes: 6 additions & 8 deletions src/viam/components/camera/camera.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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`
Expand All @@ -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
"""
...

Expand Down
21 changes: 7 additions & 14 deletions src/viam/components/camera/client.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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):
Expand All @@ -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))
Expand Down
2 changes: 1 addition & 1 deletion src/viam/components/camera/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
18 changes: 10 additions & 8 deletions tests/test_camera.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 35c64cb

Please sign in to comment.