viam-sdk 0.45.2__py3-none-win_amd64.whl

viam/media/video.py

@@ -0,0 +1,217 @@
+from array import array
+from enum import Enum
+from typing import List, Optional, Tuple
+
+from typing_extensions import Self
+
+from viam.errors import NotSupportedError
+from viam.proto.component.camera import Format
+
+from .viam_rgba import RGBA_HEADER_LENGTH, RGBA_MAGIC_NUMBER
+
+
+class CameraMimeType(str, Enum):
+    VIAM_RGBA = "image/vnd.viam.rgba"
+    VIAM_RAW_DEPTH = "image/vnd.viam.dep"
+    JPEG = "image/jpeg"
+    PNG = "image/png"
+    PCD = "pointcloud/pcd"
+
+    @classmethod
+    def from_string(cls, value: str) -> Self:
+        """Return the mimetype from a string.
+
+        Args:
+            value (str): The mimetype as a string
+
+        Returns:
+            Self: The mimetype
+        """
+        value_mime = value[:-5] if value.endswith("+lazy") else value  # ViamImage lazy encodes by default
+        return cls(value_mime)
+
+    @classmethod
+    def from_proto(cls, format: Format.ValueType) -> "CameraMimeType":
+        """Returns the mimetype from a proto enum.
+
+        Args:
+            format (Format.ValueType): The mimetype in a proto enum.
+
+        Returns:
+            Self: The mimetype.
+        """
+        mimetypes = {
+            Format.FORMAT_RAW_RGBA: CameraMimeType.VIAM_RGBA,
+            Format.FORMAT_RAW_DEPTH: CameraMimeType.VIAM_RAW_DEPTH,
+            Format.FORMAT_JPEG: CameraMimeType.JPEG,
+            Format.FORMAT_PNG: CameraMimeType.PNG,
+        }
+        return mimetypes.get(format, CameraMimeType.JPEG)
+
+    def to_proto(self) -> Format.ValueType:
+        """Returns the mimetype in a proto enum.
+
+        Returns:
+            Format.ValueType: The mimetype in a proto enum.
+        """
+        formats = {
+            self.VIAM_RGBA: Format.FORMAT_RAW_RGBA,
+            self.VIAM_RAW_DEPTH: Format.FORMAT_RAW_DEPTH,
+            self.JPEG: Format.FORMAT_JPEG,
+            self.PNG: Format.FORMAT_PNG,
+        }
+        return formats.get(self, Format.FORMAT_UNSPECIFIED)
+
+
+class ViamImage:
+    """A native implementation of an image.
+
+    Provides the raw data and the mime type.
+    """
+
+    _data: bytes
+    _mime_type: CameraMimeType
+    _height: Optional[int] = None
+    _width: Optional[int] = None
+
+    def __init__(self, data: bytes, mime_type: CameraMimeType) -> None:
+        self._data = data
+        self._mime_type = mime_type
+        self._width, self._height = _getDimensions(data, mime_type)
+
+    @property
+    def data(self) -> bytes:
+        """The raw bytes of the image"""
+        return self._data
+
+    @property
+    def mime_type(self) -> CameraMimeType:
+        """The mime type of the image"""
+        return self._mime_type
+
+    @property
+    def width(self) -> Optional[int]:
+        """The width of the image"""
+        return self._width
+
+    @property
+    def height(self) -> Optional[int]:
+        """The height of the image"""
+        return self._height
+
+    def bytes_to_depth_array(self) -> List[List[int]]:
+        """
+        Decode the data of an image that has the custom depth MIME type ``image/vnd.viam.dep`` into a standard representation.
+
+        Raises:
+            NotSupportedError: Raised if the image is not of MIME type `image/vnd.viam.dep`.
+
+        Returns:
+            List[List[int]]: The standard representation of the image.
+        """
+        if self.mime_type != CameraMimeType.VIAM_RAW_DEPTH:
+            raise NotSupportedError("Type must be `image/vnd.viam.dep` to use bytes_to_depth_array()")
+
+        self._width = int.from_bytes(self.data[8:16], "big")
+        self._height = int.from_bytes(self.data[16:24], "big")
+        depth_arr = array("H", self.data[24:])
+        depth_arr.byteswap()
+
+        depth_arr_2d = [[depth_arr[row * self._width + col] for col in range(self._width)] for row in range(self._height)]
+        return depth_arr_2d
+
+
+class NamedImage(ViamImage):
+    """An implementation of ViamImage that contains a name attribute."""
+
+    name: str
+    """The name of the image
+    """
+
+    def __init__(self, name: str, data: bytes, mime_type: CameraMimeType) -> None:
+        self.name = name
+        super().__init__(data, mime_type)
+
+
+def _getDimensions(image: bytes, mime_type: CameraMimeType) -> Tuple[Optional[int], Optional[int]]:
+    try:
+        if mime_type == CameraMimeType.JPEG:
+            return _getDimensionsFromJPEG(image)
+        if mime_type == CameraMimeType.PNG:
+            return _getDimensionsFromPNG(image)
+        if mime_type == CameraMimeType.VIAM_RGBA:
+            return _getDimensionsFromRGBA(image)
+    except ValueError:
+        return (None, None)
+    return (None, None)
+
+
+def _getDimensionsFromJPEG(image: bytes) -> Tuple[int, int]:
+    # JPEG Specification: https://www.w3.org/Graphics/JPEG/itu-t81.pdf
+    # Specification for markers: Table B.1
+
+    offset = 0
+    while offset < len(image):
+        while image[offset] == 0xFF:
+            # Skip all 0xFF bytes
+            offset += 1
+
+        marker = image[offset]
+        offset += 1
+        if marker == 0x01:
+            # Temporary/private use marker
+            offset += 1
+            continue
+        if marker in range(0xD0, 0xD7):
+            # Restart (RST) marker
+            offset += 1
+            continue
+        if marker == 0xD8:
+            # Start of image (SOI) marker
+            offset += 1
+            continue
+        if marker == 0xD9:
+            # End of image (EOI) marker
+            break
+
+        length = int.from_bytes(image[offset : offset + 1], byteorder="big")  # length of section
+        if marker == 0xC0 or marker == 0xC2:
+            height = int.from_bytes(image[offset + 3 : offset + 5], byteorder="big")
+            width = int.from_bytes(image[offset + 5 : offset + 7], byteorder="big")
+            return (width, height)
+
+        offset += length
+
+    raise ValueError("Invalid JPEG: Could not extract dimensions")
+
+
+def _getDimensionsFromPNG(image: bytes) -> Tuple[int, int]:
+    # PNG Specification: https://www.w3.org/TR/png/
+
+    # PNG will always start with this signature
+    signature = image[:8]
+    if signature != [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]:
+        ValueError("Invalid PNG: Invalid signature")
+
+    header = image[12:24]
+    chunk_type = header[:4].decode()
+    if chunk_type != "IHDR":
+        ValueError("Invalid PNG: Invalid headers")
+
+    width = int.from_bytes(header[4:8], byteorder="big")
+    height = int.from_bytes(header[8:], byteorder="big")
+    return (width, height)
+
+
+def _getDimensionsFromRGBA(image: bytes) -> Tuple[int, int]:
+    # Viam RGBA header comes in 3 4-byte chunks:
+    # * Magic Number/Signature
+    # * Width
+    # * Height
+    header = image[:RGBA_HEADER_LENGTH]
+    if header[:4] != RGBA_MAGIC_NUMBER:
+        raise ValueError("Invalid Viam RGBA: Invalid headers")
+
+    width = int.from_bytes(header[4:8], byteorder="big")
+    height = int.from_bytes(header[8:], byteorder="big")
+    return (width, height)

viam/module/__init__.py

@@ -0,0 +1,5 @@
+from .module import Module
+
+__all__ = [
+    "Module",
+]

viam/module/module.py

@@ -0,0 +1,281 @@
+import argparse
+import io
+import logging as pylogging
+import os
+import sys
+from inspect import iscoroutinefunction
+from threading import Lock
+from typing import List, Mapping, Optional, Sequence, Tuple
+
+from grpclib.metadata import Deadline
+from grpclib.utils import _service_name
+from typing_extensions import Self
+
+from viam import logging
+from viam.errors import ResourceNotFoundError, ValidationError
+from viam.logging import update_log_level
+from viam.proto.app.robot import ComponentConfig
+from viam.proto.module import (
+    AddResourceRequest,
+    HandlerDefinition,
+    HandlerMap,
+    ReadyRequest,
+    ReadyResponse,
+    ReconfigureResourceRequest,
+    RemoveResourceRequest,
+    ValidateConfigRequest,
+    ValidateConfigResponse,
+)
+from viam.proto.robot import ResourceRPCSubtype
+from viam.resource.base import ResourceBase
+from viam.resource.registry import Registry
+from viam.resource.types import API, RESOURCE_TYPE_COMPONENT, RESOURCE_TYPE_SERVICE, Model, ResourceName, resource_name_from_string
+from viam.robot.client import RobotClient
+from viam.rpc.dial import DialOptions
+from viam.rpc.server import Server
+
+from .service import ModuleRPCService
+from .types import Reconfigurable, Stoppable
+
+
+def _parse_module_args() -> argparse.Namespace:
+    """
+    Parse command-line args. Used by the various `Module` entrypoints.
+    """
+    p = argparse.ArgumentParser(description="Start this viam python module")
+    p.add_argument("socket_path", help="path where this module will serve a unix socket")
+    p.add_argument("--log-level", type=lambda name: pylogging._nameToLevel[name.upper()], default=logging.INFO)
+    return p.parse_args()
+
+
+class Module:
+    _address: str
+    _parent_address: Optional[str] = None
+    _ready: bool
+    _log_level: int
+    _lock: Lock
+    parent: Optional[RobotClient] = None
+    server: Server
+    logger: pylogging.Logger
+
+    @classmethod
+    def from_args(cls) -> Self:
+        """Create a new Module with the args provided in the command line. The first argument after the command must be
+        the socket path. If the second argument after the command is "--log-level=debug", the Module's logger will be
+        DEBUG level. Otherwise, it will be INFO level. See LogLevel documentation in the RDK for more information on how
+        to start modules with a "log-level" commandline argument.
+
+        Raises:
+            Exception: If there is no socket path provided in the command line argument
+
+        Returns:
+            Module: a new Module instance
+        """
+        args = _parse_module_args()
+        return cls(args.socket_path, log_level=args.log_level)
+
+    @classmethod
+    async def run_with_models(cls, *models: ResourceBase):
+        """
+        Module entrypoint that takes a list of ResourceBase implementations.
+        In most cases you'll want to use run_from_registry instead (see below).
+        """
+        module = cls.from_args()
+        for model in models:
+            if not hasattr(model, "MODEL"):
+                raise TypeError(f"missing MODEL field on {model}. Resource implementations must define MODEL")
+            module.add_model_from_registry(model.API, model.MODEL)  # pyright: ignore [reportAttributeAccessIssue]
+        await module.start()
+
+    @classmethod
+    async def run_from_registry(cls):
+        """
+        Module entrypoint that automatically includes all the resources you've created in your program.
+
+        Example:
+
+        if __name__ == '__main__':
+            asyncio.run(Module.run_from_registry())
+
+        Full example at examples/easy_resource/main.py.
+        """
+        module = cls.from_args()
+        for key in Registry.REGISTERED_RESOURCE_CREATORS().keys():
+            module.add_model_from_registry(*key.split("/"))  # pyright: ignore [reportArgumentType]
+        await module.start()
+
+    def __init__(self, address: str, *, log_level: int = logging.INFO) -> None:
+        # When a module is launched by viam-server, its stdout is not connected to a tty.  In
+        # response, python disables line buffering, which prevents `print` statements from being
+        # immediately flushed to viam-server. This behavior can be confusing, interfere with
+        # debugging, and is non-standard when compared to other languages.  Here, stdout and stderr
+        # are reconfigured to immediately flush.
+        if isinstance(sys.stdout, io.TextIOWrapper):
+            sys.stdout.reconfigure(line_buffering=True)
+        if isinstance(sys.stderr, io.TextIOWrapper):
+            sys.stderr.reconfigure(line_buffering=True)
+        self._address = address
+        self.server = Server(resources=[], module_service=ModuleRPCService(self))
+        self._log_level = log_level
+
+        module_name = os.environ.get("VIAM_MODULE_NAME")
+        # this can happen if the user is running an old version of viam-server that doesn't set `VIAM_MODULE_NAME`
+        if module_name is None:
+            module_name = __name__
+        self.logger = logging.getLogger(module_name)
+        self._ready = True
+        self._lock = Lock()
+
+    async def _connect_to_parent(self):
+        if self.parent is None:
+            if self._parent_address is None:
+                raise ValueError("Parent address not found")
+            self.parent = await RobotClient.at_address(
+                self._parent_address,
+                RobotClient.Options(
+                    dial_options=DialOptions(disable_webrtc=True, insecure=True),
+                    log_level=self._log_level,
+                ),
+            )
+            self.logger.debug("Starting module logging")
+            logging.setParent(self.parent)
+
+    async def _get_resource(self, name: ResourceName) -> ResourceBase:
+        await self._connect_to_parent()
+        assert self.parent is not None
+        await self.parent.refresh()
+        if name.type == RESOURCE_TYPE_COMPONENT:
+            return self.parent.get_component(name)
+        elif name.type == RESOURCE_TYPE_SERVICE:
+            return self.parent.get_service(name)
+        raise ValueError("Dependency does not describe a component nor a service")
+
+    async def _get_dependencies(self, dependencies: Sequence[str]) -> Mapping[ResourceName, ResourceBase]:
+        deps: Mapping[ResourceName, ResourceBase] = {}
+        for dep in dependencies:
+            rn = resource_name_from_string(dep)
+            deps[rn] = await self._get_resource(rn)
+        return deps
+
+    async def start(self):
+        """Start the module service and gRPC server"""
+        try:
+            await self.server.serve(log_level=self._log_level, path=self._address)
+        finally:
+            await self.stop()
+
+    async def stop(self):
+        """Stop the module service and gRPC server"""
+        self.logger.debug("Shutting down module")
+        try:
+            logging.shutdown()
+            if self.parent is not None:
+                await self.parent.close()
+        except Exception as e:
+            self.logger.error("Encountered error while shutting down module", exc_info=e)
+
+    def set_ready(self, ready: bool):
+        """Set the module's ready state. The module automatically sets to READY on load. Setting to False can be useful
+        in instances where the module is not instantly ready (for example waiting on hardware)
+
+        Args:
+            ready (bool): Whether the module is ready
+        """
+        with self._lock:
+            self._ready = ready
+
+    async def add_resource(self, request: AddResourceRequest, *, deadline: Optional[Deadline] = None):
+        dependencies = await self._get_dependencies(request.dependencies)
+        config: ComponentConfig = request.config
+        api = API.from_string(config.api)
+        model = Model.from_string(config.model, ignore_errors=True)
+        creator = Registry.lookup_resource_creator(api, model)
+        resource = creator(config, dependencies)
+        if deadline is not None and deadline.time_remaining() <= 0:
+            raise TimeoutError("Deadline expired")
+        update_log_level(resource.logger, config.log_configuration.level.upper())
+        self.server.register(resource)
+
+    async def reconfigure_resource(self, request: ReconfigureResourceRequest):
+        dependencies = await self._get_dependencies(request.dependencies)
+        config: ComponentConfig = request.config
+        api = API.from_string(config.api)
+        name = config.name
+        rn = ResourceName(namespace=api.namespace, type=api.resource_type, subtype=api.resource_subtype, name=name)
+        resource = self.server.get_resource(ResourceBase, rn)
+        if isinstance(resource, Reconfigurable):
+            resource.reconfigure(config, dependencies)
+        else:
+            if isinstance(resource, Stoppable):
+                if iscoroutinefunction(resource.stop):
+                    await resource.stop()
+                else:
+                    resource.stop()
+            add_request = AddResourceRequest(config=request.config, dependencies=request.dependencies)
+            await self.server.remove_resource(rn)
+            await self.add_resource(add_request)
+
+    async def remove_resource(self, request: RemoveResourceRequest):
+        rn = resource_name_from_string(request.name)
+        resource = self.server.get_resource(ResourceBase, rn)
+        if isinstance(resource, Stoppable):
+            if iscoroutinefunction(resource.stop):
+                await resource.stop()
+            else:
+                resource.stop()
+        await self.server.remove_resource(rn)
+
+    async def ready(self, request: ReadyRequest) -> ReadyResponse:
+        self._parent_address = request.parent_address
+        await self._connect_to_parent()
+
+        svcname_to_models: Mapping[Tuple[str, API], List[Model]] = {}
+        for api_model_str in Registry.REGISTERED_RESOURCE_CREATORS().keys():
+            api_str, model_str = api_model_str.split("/")
+            api = API.from_string(api_str)
+            model = Model.from_string(model_str)
+
+            registration = Registry.lookup_api(api)
+            service = registration.rpc_service(self.server)
+            service_name = _service_name(service)
+
+            models = svcname_to_models.get((service_name, api), [])
+            models.append(model)
+            svcname_to_models[(service_name, api)] = models
+
+        handlers: List[HandlerDefinition] = []
+        for key, value in svcname_to_models.items():
+            svc_name, api = key
+            rpc_subtype = ResourceRPCSubtype(
+                subtype=ResourceName(
+                    namespace=api.namespace,
+                    type=api.resource_type,
+                    subtype=api.resource_subtype,
+                    name="",
+                ),
+                proto_service=svc_name,
+            )
+            handler_def = HandlerDefinition(subtype=rpc_subtype, models=[str(model) for model in value])
+            handlers.append(handler_def)
+
+        return ReadyResponse(ready=self._ready, handlermap=HandlerMap(handlers=handlers))
+
+    def add_model_from_registry(self, api: API, model: Model):
+        """Add a pre-registered model to this Module"""
+
+        # All we need to do is double check that the model has already been registered
+        try:
+            Registry.lookup_resource_creator(api, model)
+        except ResourceNotFoundError:
+            raise ValueError(f"Cannot add model because it has not been registered. API: {api}. Model: {model}")
+
+    async def validate_config(self, request: ValidateConfigRequest) -> ValidateConfigResponse:
+        config: ComponentConfig = request.config
+        api = API.from_string(config.api)
+        model = Model.from_string(config.model)
+        validator = Registry.lookup_validator(api, model)
+        try:
+            dependencies = validator(config)
+            return ValidateConfigResponse(dependencies=dependencies)
+        except Exception as e:
+            raise ValidationError(f"{type(Exception)}: {e}").grpc_error

viam/module/service.py

@@ -0,0 +1,66 @@
+from typing import TYPE_CHECKING
+
+from grpclib import Status
+from grpclib.server import Stream
+
+from viam.errors import ViamGRPCError
+from viam.proto.module import (
+    AddResourceRequest,
+    AddResourceResponse,
+    ModuleServiceBase,
+    ReadyRequest,
+    ReadyResponse,
+    ReconfigureResourceRequest,
+    ReconfigureResourceResponse,
+    RemoveResourceRequest,
+    RemoveResourceResponse,
+    ValidateConfigRequest,
+    ValidateConfigResponse,
+)
+
+if TYPE_CHECKING:
+    from .module import Module
+
+
+class ModuleRPCService(ModuleServiceBase):
+    _module: "Module"
+
+    def __init__(self, module: "Module") -> None:
+        self._module = module
+
+    async def AddResource(self, stream: Stream[AddResourceRequest, AddResourceResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        try:
+            await self._module.add_resource(request, deadline=stream.deadline)
+        except TimeoutError:
+            raise ViamGRPCError(
+                message="Timeout while adding resource",
+                grpc_code=Status.DEADLINE_EXCEEDED,
+            )
+        await stream.send_message(AddResourceResponse())
+
+    async def ReconfigureResource(self, stream: Stream[ReconfigureResourceRequest, ReconfigureResourceResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        await self._module.reconfigure_resource(request)
+        await stream.send_message(ReconfigureResourceResponse())
+
+    async def RemoveResource(self, stream: Stream[RemoveResourceRequest, RemoveResourceResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        await self._module.remove_resource(request)
+        await stream.send_message(RemoveResourceResponse())
+
+    async def Ready(self, stream: Stream[ReadyRequest, ReadyResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        response = await self._module.ready(request)
+        await stream.send_message(response)
+
+    async def ValidateConfig(self, stream: Stream[ValidateConfigRequest, ValidateConfigResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        response = await self._module.validate_config(request)
+        if response is not None:
+            await stream.send_message(response)

viam/module/types.py

@@ -0,0 +1,23 @@
+from typing import Any, Mapping, Optional, Protocol, runtime_checkable
+
+from viam.proto.app.robot import ComponentConfig
+from viam.proto.common import ResourceName
+from viam.resource.base import ResourceBase
+
+
+@runtime_checkable
+class Reconfigurable(Protocol):
+    """The Reconfigurable protocol defines the requirements for making a resource Reconfigurable"""
+
+    def reconfigure(self, config: ComponentConfig, dependencies: Mapping[ResourceName, ResourceBase]): ...
+
+
+@runtime_checkable
+class Stoppable(Protocol):
+    """
+    The Stoppable protocol defines the requirements for making a resource Stoppable.
+
+    All resources that physically move should be Stoppable.
+    """
+
+    def stop(self, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None, **kwargs): ...