diff --git a/anki_vector/camera.py b/anki_vector/camera.py index f0845c6..b64406a 100644 --- a/anki_vector/camera.py +++ b/anki_vector/camera.py @@ -29,15 +29,12 @@ "CameraComponent", "CameraConfig", "CameraImage"] import asyncio -from concurrent.futures import CancelledError +from concurrent import futures import io -import time import sys +import time -from . import annotate, connection, util -from .events import Events -from .exceptions import VectorCameraFeedException, VectorCameraImageCaptureException -from .messaging import protocol +import grpc try: import numpy as np @@ -49,6 +46,16 @@ except ImportError: sys.exit("Cannot import from PIL: Do `pip3 install --user Pillow` to install") +from . import annotate, connection, util +from .events import Events +from .exceptions import (VectorCameraFeedException, + VectorCameraImageCaptureException, + VectorTimeoutException) +from .messaging import protocol + + +_CAMERA_RPC_TIMEOUT = 10.0 + def _convert_to_pillow_image(image_data: bytes) -> Image.Image: """Convert raw image bytes to a Pillow Image.""" @@ -278,6 +285,7 @@ def __init__(self, robot): self._latest_image: CameraImage = None self._latest_image_id: int = None self._camera_feed_task: asyncio.Task = None + self._camera_feed_stream = None self._enabled = False self._config = None # type CameraConfig self._gain = 0.0 @@ -414,42 +422,50 @@ def init_camera_feed(self) -> None: self._enabled = True self._camera_feed_task = self.conn.loop.create_task(self._request_and_handle_images()) - def close_camera_feed(self) -> None: - """Cancel camera feed task.""" + def close_camera_feed(self, timeout: float = _CAMERA_RPC_TIMEOUT) -> None: + """Cancel the camera feed task within a bounded time. + + :param timeout: Number of seconds to wait for the stream to close. + """ + if timeout <= 0: + raise ValueError('timeout must be greater than zero') if self._camera_feed_task: + deadline = time.monotonic() + timeout self._enabled = False - self._camera_feed_task.cancel() + if self._camera_feed_stream is not None: + self._camera_feed_stream.cancel() + else: + self._camera_feed_task.cancel() future = self.conn.run_coroutine(self._camera_feed_task) try: - future.result() - except CancelledError: - self.logger.debug('Camera feed task was cancelled. This is expected during disconnection.') - # wait for streaming to end, up to 10 seconds - iterations = 0 - max_iterations = 100 - while self.image_streaming_enabled(): - time.sleep(0.1) - iterations += 1 - if iterations > max_iterations: - # leave loop, even if streaming is still enabled - # because other SDK functions will still work and - # the RPC should have had enough time to finish - # which means we _should_ be in a good state. - self.logger.info('Camera Feed closed, but streaming on' - ' robot remained enabled. This is unexpected.') + future.result(timeout=timeout) + except futures.CancelledError: + self.logger.debug( + 'Camera feed task was cancelled. This is expected during ' + 'disconnection.') + except futures.TimeoutError as exc: + future.cancel() + raise VectorTimeoutException(None) from exc + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise VectorTimeoutException(None) + if not self.image_streaming_enabled(timeout=min(1.0, remaining)): break + time.sleep(min(0.1, remaining)) self._camera_feed_task = None - async def _image_streaming_enabled(self) -> bool: + async def _image_streaming_enabled(self, timeout: float) -> bool: """request streaming enabled status from the robot""" request = protocol.IsImageStreamingEnabledRequest() - response = await self.conn.grpc_interface.IsImageStreamingEnabled(request) + response = await self.conn.grpc_interface.IsImageStreamingEnabled( + request, timeout=timeout) enabled = False if response: enabled = response.is_image_streaming_enabled return enabled - def image_streaming_enabled(self) -> bool: + def image_streaming_enabled(self, timeout: float = _CAMERA_RPC_TIMEOUT) -> bool: """True if image streaming is enabled on the robot .. testcode:: @@ -461,9 +477,17 @@ def image_streaming_enabled(self) -> bool: print("Robot is streaming video") else: print("Robot is not streaming video") + + :param timeout: Number of seconds to wait for the status RPC. """ - future = self.conn.run_coroutine(self._image_streaming_enabled()) - return future.result() + if timeout <= 0: + raise ValueError('timeout must be greater than zero') + future = self.conn.run_coroutine(self._image_streaming_enabled(timeout)) + try: + return future.result(timeout=timeout) + except futures.TimeoutError as exc: + future.cancel() + raise VectorTimeoutException(None) from exc def _unpack_image(self, msg: protocol.CameraFeedResponse) -> None: """Processes raw data from the robot into a more useful image structure.""" @@ -484,20 +508,36 @@ def _unpack_image(self, msg: protocol.CameraFeedResponse) -> None: async def _request_and_handle_images(self) -> None: """Queries and listens for camera feed events from the robot. Received events are parsed by a helper function.""" + stream = None try: req = protocol.CameraFeedRequest() - async for evt in self.grpc_interface.CameraFeed(req): + stream = self.grpc_interface.CameraFeed(req) + self._camera_feed_stream = stream + async for evt in stream: # If the camera feed is disabled after stream is setup, exit the stream # (the camera feed on the robot is disabled internally on stream exit) if not self._enabled: self.logger.warning('Camera feed has been disabled. Enable the feed to start/continue receiving camera feed data') return self._unpack_image(evt) - except CancelledError: - self.logger.debug('Camera feed task was cancelled. This is expected during disconnection.') + except asyncio.CancelledError: + self.logger.debug( + 'Camera feed task was cancelled. This is expected during ' + 'disconnection.') + raise + except grpc.RpcError as exc: + if exc.code() != grpc.StatusCode.CANCELLED: + raise + finally: + self._camera_feed_stream = None + if stream is not None: + stream.cancel() @connection.on_connection_thread() - async def capture_single_image(self, enable_high_resolution: bool = False) -> CameraImage: + async def capture_single_image( + self, + enable_high_resolution: bool = False, + timeout: float = _CAMERA_RPC_TIMEOUT) -> CameraImage: """Request to capture a single image from the robot's camera. This call requests the robot to capture an image and returns the @@ -516,14 +556,17 @@ async def capture_single_image(self, enable_high_resolution: bool = False) -> Ca :param enable_high_resolution: Enable/disable request for high resolution images. The default resolution is 640x360, while the high resolution is 1280x720. + :param timeout: Number of seconds to wait for the gRPC response. """ + if timeout <= 0: + raise ValueError('timeout must be greater than zero') if self._enabled: self.logger.warning('Camera feed is enabled. Receiving image from the feed at default resolution.') return self._latest_image if enable_high_resolution: self.logger.warning('Capturing a high resolution (1280*720) image. Image events for this frame need to be scaled.') req = protocol.CaptureSingleImageRequest(enable_high_resolution=enable_high_resolution) - res = await self.grpc_interface.CaptureSingleImage(req) + res = await self.grpc_interface.CaptureSingleImage(req, timeout=timeout) if res and res.data: image = _convert_to_pillow_image(res.data) return CameraImage(image, self._image_annotator, res.image_id)