Files

148 lines
4.6 KiB
Python

# -*- coding: utf-8 -*-
# Copyright (C) 2018-2025 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import numpy as np
from functools import singledispatchmethod
from collections.abc import Iterator, Mapping, KeysView, ItemsView, ValuesView
from typing import Union, Optional
from openvino._pyopenvino import Tensor, ConstOutput
from openvino._pyopenvino import InferRequest as InferRequestBase
def tensor_from_file(path: str) -> Tensor:
"""Create Tensor from file. Data will be read with dtype of unit8."""
return Tensor(np.fromfile(path, dtype=np.uint8)) # type: ignore
class _InferRequestWrapper(InferRequestBase):
"""InferRequest class with internal memory."""
def __init__(self, other: InferRequestBase) -> None:
# Private memeber to store newly created shared memory data
self._inputs_data = None
super().__init__(other)
def _is_single_input(self) -> bool:
return len(self.input_tensors) == 1
class OVDict(Mapping):
"""Custom OpenVINO dictionary with inference results.
This class is a dict-like object. It provides possibility to
address data tensors with three key types:
* `openvino.ConstOutput` - port of the output
* `int` - index of the output
* `str` - names of the output
This class follows `frozenset`/`tuple` concept of immutability.
It is prohibited to assign new items or edit them.
To revert to the previous behavior use `to_dict` method which
return shallow copy of underlaying dictionary.
Note: It removes addressing feature! New dictionary keeps
only `ConstOutput` keys.
If a tuple returns value is needed, use `to_tuple` method which
converts values to the tuple.
:Example:
.. code-block:: python
# Reverts to the previous behavior of the native dict
result = request.infer(inputs).to_dict()
# or alternatively:
result = dict(request.infer(inputs))
.. code-block:: python
# To dispatch outputs of multi-ouput inference:
out1, out2, out3, _ = request.infer(inputs).values()
# or alternatively:
out1, out2, out3, _ = request.infer(inputs).to_tuple()
"""
def __init__(self, _dict: dict[ConstOutput, np.ndarray]) -> None:
self._dict = _dict
self._names: Optional[dict[ConstOutput, set[str]]] = None
def __iter__(self) -> Iterator:
return self._dict.__iter__()
def __len__(self) -> int:
return len(self._dict)
def __repr__(self) -> str:
return self._dict.__repr__()
def __get_names(self) -> dict[ConstOutput, set[str]]:
"""Return names of every output key.
Insert empty set if key has no name.
"""
return {key: key.get_names() for key in self._dict.keys()}
def __get_key(self, index: int) -> ConstOutput:
return list(self._dict.keys())[index]
@singledispatchmethod
def __getitem_impl(self, key: Union[ConstOutput, int, str]) -> np.ndarray:
raise TypeError(f"Unknown key type: {type(key)}")
@__getitem_impl.register
def _(self, key: ConstOutput) -> np.ndarray:
return self._dict[key]
@__getitem_impl.register
def _(self, key: int) -> np.ndarray:
try:
return self._dict[self.__get_key(key)]
except IndexError:
raise KeyError(key)
@__getitem_impl.register
def _(self, key: str) -> np.ndarray:
if self._names is None:
self._names = self.__get_names()
for port, port_names in self._names.items():
if key in port_names:
return self._dict[port]
raise KeyError(key)
def __getitem__(self, key: Union[ConstOutput, int, str]) -> np.ndarray:
return self.__getitem_impl(key)
def keys(self) -> KeysView[ConstOutput]:
return self._dict.keys()
def values(self) -> ValuesView[np.ndarray]:
return self._dict.values()
def items(self) -> ItemsView[ConstOutput, np.ndarray]:
return self._dict.items()
def names(self) -> tuple[set[str], ...]:
"""Return names of every output key.
Insert empty set if key has no name.
"""
if self._names is None:
self._names = self.__get_names()
return tuple(self._names.values())
def to_dict(self) -> dict[ConstOutput, np.ndarray]:
"""Return underlaying native dictionary.
Function performs shallow copy, thus any modifications to
returned values may affect this class as well.
"""
return self._dict
def to_tuple(self) -> tuple:
"""Convert values of this dictionary to a tuple."""
return tuple(self._dict.values())