RSDK-4895: Make numpy an extra dependency (#429)

hexbabe · web-flow · commit d4cf92d82585 · 2023-09-14T16:49:14.000-04:00
Co-authored-by: hexbabe &lt;sean.yu@viam.com&gt;
diff --git a/Makefile b/Makefile
@@ -70,5 +70,5 @@ package: clean buf better_imports format test
 	@echo "TODO: Create pip-installable package"
 
 install:
-	poetry install
+	poetry install --all-extras
 	sh etc/postinstall.sh
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -17,7 +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"
+  numpy = { version = ">=1.21", optional = true }
 
   [tool.poetry.group.dev.dependencies]
   pytest = "^7.4.0"
@@ -60,3 +60,6 @@ line_length = 140
 [build-system]
 requires = [ "poetry-core>=1.0.0" ]
 build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.extras]
+mlmodel = ["numpy"]
diff --git a/src/viam/services/mlmodel/__init__.py b/src/viam/services/mlmodel/__init__.py
@@ -1,3 +1,17 @@
+try:
+    import numpy
+except ImportError:
+    import warnings
+
+    warnings.warn(
+        (
+            """MLModel support in the Viam Python SDK requires the installation of an
+additional dependency: numpy.  Update your package using the extra [mlmodel]
+e.g. `pip install viam-sdk[mlmodel]` or the equivalent update in your dependency manager."""
+        ),
+    )
+    raise
+
 from viam.proto.service.mlmodel import File, LabelType, Metadata, TensorInfo
 from viam.resource.registry import Registry, ResourceRegistration
 
diff --git a/src/viam/services/mlmodel/client.py b/src/viam/services/mlmodel/client.py
@@ -1,12 +1,12 @@
-from typing import Dict, Mapping, Optional
-
 from grpclib.client import Channel
 from numpy.typing import NDArray
+from typing import Dict, Mapping, Optional
 
 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, flat_tensors_to_ndarrays, ndarrays_to_flat_tensors, struct_to_dict
+from viam.services.mlmodel.utils import flat_tensors_to_ndarrays, ndarrays_to_flat_tensors
+from viam.utils import ValueTypes, dict_to_struct, struct_to_dict
 
 from .mlmodel import Metadata, MLModel
 
diff --git a/src/viam/services/mlmodel/service.py b/src/viam/services/mlmodel/service.py
@@ -2,7 +2,7 @@
 
 from viam.proto.service.mlmodel import InferRequest, InferResponse, MetadataRequest, MetadataResponse, MLModelServiceBase
 from viam.resource.rpc_service_base import ResourceRPCServiceBase
-from viam.utils import flat_tensors_to_ndarrays, ndarrays_to_flat_tensors
+from viam.services.mlmodel.utils import flat_tensors_to_ndarrays, ndarrays_to_flat_tensors
 
 from .mlmodel import MLModel
 
diff --git a/src/viam/services/mlmodel/utils.py b/src/viam/services/mlmodel/utils.py
@@ -0,0 +1,90 @@
+import numpy as np
+from numpy.typing import NDArray
+from typing import Dict
+
+from viam.proto.service.mlmodel import (
+    FlatTensors,
+    FlatTensor,
+    FlatTensorDataDouble,
+    FlatTensorDataFloat,
+    FlatTensorDataInt16,
+    FlatTensorDataInt32,
+    FlatTensorDataInt64,
+    FlatTensorDataInt8,
+    FlatTensorDataUInt16,
+    FlatTensorDataUInt32,
+    FlatTensorDataUInt64,
+    FlatTensorDataUInt8,
+)
+
+
+def flat_tensors_to_ndarrays(flat_tensors: FlatTensors) -> Dict[str, NDArray]:
+    property_name_to_dtype = {
+        "float_tensor": np.float32,
+        "double_tensor": np.float64,
+        "int8_tensor": np.int8,
+        "int16_tensor": np.int16,
+        "int32_tensor": np.int32,
+        "int64_tensor": np.int64,
+        "uint8_tensor": np.uint8,
+        "uint16_tensor": np.uint16,
+        "uint32_tensor": np.uint32,
+        "uint64_tensor": np.uint64,
+    }
+
+    def make_ndarray(flat_data, dtype, shape):
+        """Takes flat data (protobuf RepeatedScalarFieldContainer | bytes) to output an ndarray
+        of appropriate dtype and shape"""
+        make_array = np.frombuffer if dtype == np.int8 or dtype == np.uint8 else np.array
+        return make_array(flat_data, dtype).reshape(shape)
+
+    ndarrays: Dict[str, NDArray] = dict()
+    for name, flat_tensor in flat_tensors.tensors.items():
+        property_name = flat_tensor.WhichOneof("tensor") or flat_tensor.WhichOneof(b"tensor")
+        if property_name:
+            tensor_data = getattr(flat_tensor, property_name)
+            flat_data, dtype, shape = tensor_data.data, property_name_to_dtype[property_name], flat_tensor.shape
+            ndarrays[name] = make_ndarray(flat_data, dtype, shape)
+    return ndarrays
+
+
+def ndarrays_to_flat_tensors(ndarrays: Dict[str, NDArray]) -> FlatTensors:
+    dtype_name_to_tensor_data_class = {
+        "float32": FlatTensorDataFloat,
+        "float64": FlatTensorDataDouble,
+        "int8": FlatTensorDataInt8,
+        "int16": FlatTensorDataInt16,
+        "int32": FlatTensorDataInt32,
+        "int64": FlatTensorDataInt64,
+        "uint8": FlatTensorDataUInt8,
+        "uint16": FlatTensorDataUInt16,
+        "uint32": FlatTensorDataUInt32,
+        "uint64": FlatTensorDataUInt64,
+    }
+
+    def get_tensor_data(ndarray: NDArray):
+        """Takes an ndarray and returns the corresponding tensor data class instance
+        e.g. FlatTensorDataInt8, FlatTensorDataUInt8 etc."""
+        tensor_data_class = dtype_name_to_tensor_data_class[ndarray.dtype.name]
+        data = ndarray.flatten()
+        if tensor_data_class == FlatTensorDataInt8 or tensor_data_class == FlatTensorDataUInt8:
+            data = data.tobytes()  # as per the proto, int8 and uint8 are stored as bytes
+        elif tensor_data_class == FlatTensorDataInt16 or tensor_data_class == FlatTensorDataUInt16:
+            data = data.astype(np.uint32)  # as per the proto, int16 and uint16 are stored as uint32
+        tensor_data = tensor_data_class(data=data)
+        return tensor_data
+
+    def get_tensor_data_type(ndarray: NDArray):
+        """Takes ndarray and returns a FlatTensor datatype property to be set
+        e.g. "float_tensor", "uint32_tensor" etc."""
+        if ndarray.dtype == np.float32:
+            return "float_tensor"
+        elif ndarray.dtype == np.float64:
+            return "double_tensor"
+        return f"{ndarray.dtype.name}_tensor"
+
+    tensors_mapping: Dict[str, FlatTensor] = dict()
+    for name, ndarray in ndarrays.items():
+        prop_name, prop_value = get_tensor_data_type(ndarray), get_tensor_data(ndarray)
+        tensors_mapping[name] = FlatTensor(shape=ndarray.shape, **{prop_name: prop_value})
+    return FlatTensors(tensors=tensors_mapping)
diff --git a/src/viam/utils.py b/src/viam/utils.py
@@ -1,37 +1,23 @@
 import asyncio
 import contextvars
 import functools
+
 import sys
 import threading
 from datetime import datetime
 from typing import Any, Dict, List, Mapping, Optional, SupportsBytes, SupportsFloat, Type, TypeVar, Union
 
-import numpy as np
 from google.protobuf.json_format import MessageToDict, ParseDict
 from google.protobuf.message import Message
 from google.protobuf.struct_pb2 import ListValue, Struct, Value
 from google.protobuf.timestamp_pb2 import Timestamp
-from numpy.typing import NDArray
 
 from viam.proto.common import Geometry, GeoPoint, GetGeometriesRequest, GetGeometriesResponse, Orientation, ResourceName, Vector3
-from viam.proto.service.mlmodel import (
-    FlatTensor,
-    FlatTensorDataDouble,
-    FlatTensorDataFloat,
-    FlatTensorDataInt8,
-    FlatTensorDataInt16,
-    FlatTensorDataInt32,
-    FlatTensorDataInt64,
-    FlatTensorDataUInt8,
-    FlatTensorDataUInt16,
-    FlatTensorDataUInt32,
-    FlatTensorDataUInt64,
-    FlatTensors,
-)
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry
 from viam.resource.types import Subtype, SupportsGetGeometries
 
+
 if sys.version_info >= (3, 9):
     from collections.abc import Callable
 else:
@@ -165,78 +151,6 @@ def _convert(v: ValueTypes) -> Any:
     return struct
 
 
-def flat_tensors_to_ndarrays(flat_tensors: FlatTensors) -> Dict[str, NDArray]:
-    property_name_to_dtype = {
-        "float_tensor": np.float32,
-        "double_tensor": np.float64,
-        "int8_tensor": np.int8,
-        "int16_tensor": np.int16,
-        "int32_tensor": np.int32,
-        "int64_tensor": np.int64,
-        "uint8_tensor": np.uint8,
-        "uint16_tensor": np.uint16,
-        "uint32_tensor": np.uint32,
-        "uint64_tensor": np.uint64,
-    }
-
-    def make_ndarray(flat_data, dtype, shape):
-        """Takes flat data (protobuf RepeatedScalarFieldContainer | bytes) to output an ndarray
-        of appropriate dtype and shape"""
-        make_array = np.frombuffer if dtype == np.int8 or dtype == np.uint8 else np.array
-        return make_array(flat_data, dtype).reshape(shape)
-
-    ndarrays: Dict[str, NDArray] = dict()
-    for name, flat_tensor in flat_tensors.tensors.items():
-        property_name = flat_tensor.WhichOneof("tensor") or flat_tensor.WhichOneof(b"tensor")
-        if property_name:
-            tensor_data = getattr(flat_tensor, property_name)
-            flat_data, dtype, shape = tensor_data.data, property_name_to_dtype[property_name], flat_tensor.shape
-            ndarrays[name] = make_ndarray(flat_data, dtype, shape)
-    return ndarrays
-
-
-def ndarrays_to_flat_tensors(ndarrays: Dict[str, NDArray]) -> FlatTensors:
-    dtype_name_to_tensor_data_class = {
-        "float32": FlatTensorDataFloat,
-        "float64": FlatTensorDataDouble,
-        "int8": FlatTensorDataInt8,
-        "int16": FlatTensorDataInt16,
-        "int32": FlatTensorDataInt32,
-        "int64": FlatTensorDataInt64,
-        "uint8": FlatTensorDataUInt8,
-        "uint16": FlatTensorDataUInt16,
-        "uint32": FlatTensorDataUInt32,
-        "uint64": FlatTensorDataUInt64,
-    }
-
-    def get_tensor_data(ndarray: NDArray):
-        """Takes an ndarray and returns the corresponding tensor data class instance
-        e.g. FlatTensorDataInt8, FlatTensorDataUInt8 etc."""
-        tensor_data_class = dtype_name_to_tensor_data_class[ndarray.dtype.name]
-        data = ndarray.flatten()
-        if tensor_data_class == FlatTensorDataInt8 or tensor_data_class == FlatTensorDataUInt8:
-            data = data.tobytes()  # as per the proto, int8 and uint8 are stored as bytes
-        elif tensor_data_class == FlatTensorDataInt16 or tensor_data_class == FlatTensorDataUInt16:
-            data = data.astype(np.uint32)  # as per the proto, int16 and uint16 are stored as uint32
-        tensor_data = tensor_data_class(data=data)
-        return tensor_data
-
-    def get_tensor_data_type(ndarray: NDArray):
-        """Takes ndarray and returns a FlatTensor datatype property to be set
-        e.g. "float_tensor", "uint32_tensor" etc."""
-        if ndarray.dtype == np.float32:
-            return "float_tensor"
-        elif ndarray.dtype == np.float64:
-            return "double_tensor"
-        return f"{ndarray.dtype.name}_tensor"
-
-    tensors_mapping: Dict[str, FlatTensor] = dict()
-    for name, ndarray in ndarrays.items():
-        prop_name, prop_value = get_tensor_data_type(ndarray), get_tensor_data(ndarray)
-        tensors_mapping[name] = FlatTensor(shape=ndarray.shape, **{prop_name: prop_value})
-    return FlatTensors(tensors=tensors_mapping)
-
-
 def struct_to_dict(struct: Struct) -> Dict[str, ValueTypes]:
     return {key: value_to_primitive(value) for (key, value) in struct.fields.items()}
 
diff --git a/tests/mocks/services.py b/tests/mocks/services.py
@@ -227,10 +227,11 @@
 from viam.proto.service.vision import Classification, Detection
 from viam.app.data_client import DataClient
 from viam.services.mlmodel import File, LabelType, Metadata, MLModel, TensorInfo
+from viam.services.mlmodel.utils import ndarrays_to_flat_tensors, flat_tensors_to_ndarrays
 from viam.services.navigation import Navigation
 from viam.services.slam import SLAM
 from viam.services.vision import Vision
-from viam.utils import ValueTypes, datetime_to_timestamp, dict_to_struct, struct_to_dict, ndarrays_to_flat_tensors, flat_tensors_to_ndarrays
+from viam.utils import ValueTypes, datetime_to_timestamp, dict_to_struct, struct_to_dict
 
 
 class MockVision(Vision):
diff --git a/tests/test_mlmodel.py b/tests/test_mlmodel.py
@@ -64,6 +64,9 @@ def setup_class(cls):
         cls.manager = ResourceManager([cls.mlmodel])
         cls.service = MLModelRPCService(cls.manager)
 
+    #  ignore warning about our out-of-bound int casting (i.e. uint32 -> int16) because we don't store uint32s for int16 & uint16 tensor
+    #  data > 2^16-1 in the first place (inherently they are int16, we just cast them to uint32 for the grpc message)
+    @pytest.mark.filterwarnings("ignore::DeprecationWarning")
     @pytest.mark.asyncio
     async def test_infer(self):
         async with ChannelFor([self.service]) as channel:
diff --git a/tests/test_mlmodel_utils.py b/tests/test_mlmodel_utils.py
@@ -0,0 +1,57 @@
+import builtins
+import numpy as np
+import pytest
+
+from .mocks.services import MockMLModel
+
+from viam.services.mlmodel.utils import flat_tensors_to_ndarrays, ndarrays_to_flat_tensors
+
+
+#  ignore warning about our out-of-bound int casting (i.e. uint32 -> int16) because we don't store uint32s for int16 & uint16 tensor data
+# > 2^16-1 in the first place (inherently they are int16, we just cast them to uint32 for the grpc message)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_flat_tensors_to_ndarrays():
+    output = flat_tensors_to_ndarrays(MockMLModel.INTS_FLAT_TENSORS)
+    assert len(output.keys()) == 4
+    assert all(name in output.keys() for name in ["0", "1", "2", "3"])
+    assert np.array_equal(output["0"], MockMLModel.INT8_NDARRAY)
+    assert output["0"].dtype == np.int8
+    assert np.array_equal(output["1"], MockMLModel.INT16_NDARRAY)
+    assert output["1"].dtype == np.int16
+    assert np.array_equal(output["2"], MockMLModel.INT32_NDARRAY)
+    assert output["2"].dtype == np.int32
+    assert np.array_equal(output["3"], MockMLModel.INT64_NDARRAY)
+    assert output["3"].dtype == np.int64
+
+    output = flat_tensors_to_ndarrays(MockMLModel.UINTS_FLAT_TENSORS)
+    assert len(output.keys()) == 4
+    assert all(name in output.keys() for name in ["0", "1", "2", "3"])
+    assert np.array_equal(output["0"], MockMLModel.UINT8_NDARRAY)
+    assert output["0"].dtype == np.uint8
+    assert np.array_equal(output["1"], MockMLModel.UINT16_NDARRAY)
+    assert output["1"].dtype == np.uint16
+    assert np.array_equal(output["2"], MockMLModel.UINT32_NDARRAY)
+    assert output["2"].dtype == np.uint32
+    assert np.array_equal(output["3"], MockMLModel.UINT64_NDARRAY)
+    assert output["3"].dtype == np.uint64
+
+    output = flat_tensors_to_ndarrays(MockMLModel.DOUBLE_FLOAT_TENSORS)
+    assert len(output.keys()) == 2
+    assert all(name in output.keys() for name in ["0", "1"])
+    assert np.array_equal(output["0"], MockMLModel.DOUBLE_NDARRAY)
+    assert output["0"].dtype == np.float64
+    assert np.array_equal(output["1"], MockMLModel.FLOAT_NDARRAY)
+    assert output["1"].dtype == np.float32
+
+
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_ndarrays_to_flat_tensors():
+    output = ndarrays_to_flat_tensors(MockMLModel.INTS_NDARRAYS)
+    assert len(output.tensors) == 4
+    assert all(name in output.tensors.keys() for name in ["0", "1", "2", "3"])
+    assert type(output.tensors["0"].int8_tensor.data) is builtins.bytes
+    bytes_buffer = output.tensors["0"].int8_tensor.data
+    assert np.array_equal(np.frombuffer(bytes_buffer, dtype=np.int8).reshape(output.tensors["0"].shape), MockMLModel.INT8_NDARRAY)
+    assert np.array_equal(np.array(output.tensors["1"].int16_tensor.data, dtype=np.int16), MockMLModel.INT16_NDARRAY)
+    assert np.array_equal(np.array(output.tensors["2"].int32_tensor.data, dtype=np.int32), MockMLModel.INT32_NDARRAY)
+    assert np.array_equal(np.array(output.tensors["3"].int64_tensor.data, dtype=np.int64), MockMLModel.INT64_NDARRAY)
diff --git a/tests/test_utils.py b/tests/test_utils.py
