RSDK-4197: Add support for Flat Tensors in Python SDK

examples/server/v1/services.py

@@ -1,19 +1,19 @@
 from datetime import datetime
+from numpy.typing import NDArray
 from typing import Dict, List, Optional
 from tests.mocks.services import MockMLModel, MockSLAM
 from viam.services.slam import Pose, SLAM
 from viam.services.mlmodel import Metadata, MLModel
-from viam.utils import ValueTypes
 
 
 class ExampleMLModel(MLModel):
     def __init__(self, name: str):
-        self.input_data = MockMLModel.INPUT_DATA
-        self.output_data = MockMLModel.OUTPUT_DATA
+        self.input_data = MockMLModel.EMPTY_NDARRAYS
+        self.output_data = MockMLModel.EMPTY_NDARRAYS
         self.meta = MockMLModel.META
         super().__init__(name)
 
-    async def infer(self, input_data: Dict[str, ValueTypes], *, timeout: Optional[float] = None) -> Dict[str, ValueTypes]:
+    async def infer(self, input_tensors: Dict[str, NDArray], *, timeout: Optional[float] = None) -> Dict[str, NDArray]:
         return self.output_data
 
     async def metadata(self, *, timeout: Optional[float] = None) -> Metadata:

pyproject.toml

@@ -17,6 +17,7 @@ include = ["LICENSE", "src/viam/rpc/libviam_rust_utils.*"]
   typing-extensions = "^4.7.1"
   Pillow = "^10.0.0"
   protobuf = "^4.23.4"
+  numpy = ">=1.21"
 
   [tool.poetry.group.dev.dependencies]
   pytest = "^7.4.0"

src/viam/services/mlmodel/client.py

@@ -5,7 +5,8 @@
 from viam.proto.common import DoCommandRequest, DoCommandResponse
 from viam.proto.service.mlmodel import InferRequest, InferResponse, MetadataRequest, MetadataResponse, MLModelServiceStub
 from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase
-from viam.utils import ValueTypes, dict_to_struct, struct_to_dict
+from viam.utils import ValueTypes, dict_to_struct, struct_to_dict, flat_tensors_to_ndarrays, ndarrays_to_flat_tensors
+from numpy.typing import NDArray
 
 from .mlmodel import Metadata, MLModel
 
@@ -16,10 +17,10 @@ def __init__(self, name: str, channel: Channel):
         self.client = MLModelServiceStub(channel)
         super().__init__(name)
 
-    async def infer(self, input_data: Dict[str, ValueTypes], *, timeout: Optional[float] = None) -> Dict[str, ValueTypes]:
-        request = InferRequest(name=self.name, input_data=dict_to_struct(input_data))
+    async def infer(self, input_tensors: Dict[str, NDArray], *, timeout: Optional[float] = None) -> Dict[str, NDArray]:
+        request = InferRequest(name=self.name, input_tensors=ndarrays_to_flat_tensors(input_tensors))
         response: InferResponse = await self.client.Infer(request)
-        return struct_to_dict(response.output_data)
+        return flat_tensors_to_ndarrays(response.output_tensors)
 
     async def metadata(self, *, timeout: Optional[float] = None) -> Metadata:
         request = MetadataRequest(name=self.name)
@@ -28,5 +29,5 @@ async def metadata(self, *, timeout: Optional[float] = None) -> Metadata:
 
     async def do_command(self, command: Mapping[str, ValueTypes], *, timeout: Optional[float] = None, **kwargs) -> Mapping[str, ValueTypes]:
         request = DoCommandRequest(name=self.name, command=dict_to_struct(command))
-        response: DoCommandResponse = await self.client.DoCommand(request, tiemout=timeout)
+        response: DoCommandResponse = await self.client.DoCommand(request, timeout=timeout)
         return struct_to_dict(response.result)

src/viam/services/mlmodel/mlmodel.py

@@ -1,9 +1,9 @@
 import abc
+from numpy.typing import NDArray
 from typing import Dict, Final, Optional
 
 from viam.proto.service.mlmodel import Metadata
 from viam.resource.types import RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_SERVICE, Subtype
-from viam.utils import ValueTypes
 
 from ..service_base import ServiceBase
 
@@ -20,14 +20,14 @@ class MLModel(ServiceBase):
     SUBTYPE: Final = Subtype(RESOURCE_NAMESPACE_RDK, RESOURCE_TYPE_SERVICE, "mlmodel")
 
     @abc.abstractmethod
-    async def infer(self, input_data: Dict[str, ValueTypes], *, timeout: Optional[float]) -> Dict[str, ValueTypes]:
+    async def infer(self, input_tensors: Dict[str, NDArray], *, timeout: Optional[float]) -> Dict[str, NDArray]:
         """Take an already ordered input tensor as an array, make an inference on the model, and return an output tensor map.
 
         Args:
-            input_data (Dict[str, ValueTypes]): A dictionary of input arrays/tensors as specified in the metadata
+            input_tensors (Dict[str, NDArray]): A dictionary of input flat tensors as specified in the metadata
 
         Returns:
-            Dict[str, ValueTypes]: A dictionary of output arrays/tensors as specified in the metadata
+            Dict[str, NDArray]: A dictionary of output flat tensors as specified in the metadata
         """
         ...