70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
# Copyright (C) 2020-2022 Intel Corporation
|
|
#
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from openvino.inference_engine import IECore
|
|
|
|
class ModelLoader:
|
|
def __init__(self, model, weights):
|
|
ie_core = IECore()
|
|
network = ie_core.read_network(model, weights)
|
|
self._network = network
|
|
|
|
# Initialize input blobs
|
|
self._input_info_name = None
|
|
for blob_name in network.input_info:
|
|
if len(network.input_info[blob_name].tensor_desc.dims) == 4:
|
|
self._input_blob_name = blob_name
|
|
self._input_layout = network.input_info[blob_name].tensor_desc.dims
|
|
elif len(network.input_info[blob_name].tensor_desc.dims) == 2:
|
|
self._input_info_name = blob_name
|
|
else:
|
|
raise RuntimeError(
|
|
"Unsupported {}D input layer '{}'. Only 2D and 4D input layers are supported"
|
|
.format(len(network.input_info[blob_name].tensor_desc.dims), blob_name))
|
|
|
|
# Initialize output blob
|
|
self._output_blob_name = next(iter(network.outputs))
|
|
|
|
# Load network
|
|
self._net = ie_core.load_network(network, "CPU", num_requests=2)
|
|
|
|
def _prepare_inputs(self, image, preprocessing):
|
|
image = np.array(image)
|
|
_, _, h, w = self._input_layout
|
|
if preprocessing:
|
|
image = image if image.shape[:-1] == (h, w) else cv2.resize(image, (w, h))
|
|
if len(image.shape) < 3: # grayscale image
|
|
image = image[:, :, np.newaxis]
|
|
else:
|
|
if image.shape[2] == 4: # the image has alpha channel
|
|
image = image[:, :, :3]
|
|
|
|
image = image.transpose((2, 0, 1)) # Change data layout from HWC to CHW
|
|
|
|
inputs = {self._input_blob_name: image}
|
|
if self._input_info_name:
|
|
inputs[self._input_info_name] = [h, w, 1]
|
|
return inputs
|
|
|
|
def infer(self, image, preprocessing=True):
|
|
inputs = self._prepare_inputs(image, preprocessing)
|
|
results = self._net.infer(inputs)
|
|
if len(results) == 1:
|
|
return results[self._output_blob_name].copy()
|
|
else:
|
|
return results.copy()
|
|
|
|
def async_infer(self, image, preprocessing=True, request_id=0):
|
|
inputs = self._prepare_inputs(image, preprocessing)
|
|
return self._net.start_async(request_id=request_id, inputs=inputs)
|
|
|
|
def input_size(self):
|
|
return self._input_layout[2:]
|
|
|
|
@property
|
|
def network(self):
|
|
return self._network
|