Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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
5 changes: 4 additions & 1 deletion poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 4 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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"]
14 changes: 14 additions & 0 deletions src/viam/services/mlmodel/__init__.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
6 changes: 3 additions & 3 deletions src/viam/services/mlmodel/client.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/viam/services/mlmodel/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
90 changes: 90 additions & 0 deletions src/viam/services/mlmodel/utils.py
Original file line number Diff line number Diff line change
@@ -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)
90 changes: 2 additions & 88 deletions src/viam/utils.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down Expand Up @@ -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()}

Expand Down
3 changes: 2 additions & 1 deletion tests/mocks/services.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
3 changes: 3 additions & 0 deletions tests/test_mlmodel.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
57 changes: 57 additions & 0 deletions tests/test_mlmodel_utils.py
Original file line number Diff line number Diff line change
@@ -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)
Loading