diff --git a/docs/examples/example.ipynb b/docs/examples/example.ipynb index 138db71dd..afeff7d2f 100644 --- a/docs/examples/example.ipynb +++ b/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" diff --git a/src/viam/components/camera/camera.py b/src/viam/components/camera/camera.py index b7cc0aa9f..0b40d3a66 100644 --- a/src/viam/components/camera/camera.py +++ b/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 """ ... diff --git a/src/viam/components/camera/client.py b/src/viam/components/camera/client.py index 30ed1d3c4..323253f2f 100644 --- a/src/viam/components/camera/client.py +++ b/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)) diff --git a/src/viam/components/camera/service.py b/src/viam/components/camera/service.py index eeb8e1edc..a325eb505 100644 --- a/src/viam/components/camera/service.py +++ b/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: diff --git a/tests/test_camera.py b/tests/test_camera.py index 429d24284..87fd7b8ac 100644 --- a/tests/test_camera.py +++ b/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):