You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
import cv2
|
|
import numpy as np
|
|
from .base.base import HamburgerABC
|
|
from hyperlpr3.common.tools_process import cost
|
|
|
|
|
|
def encode_images(image: np.ndarray):
|
|
image_encode = image / 255.0
|
|
if len(image_encode.shape) == 4:
|
|
image_encode = image_encode.transpose(0, 3, 1, 2)
|
|
else:
|
|
image_encode = image_encode.transpose(2, 0, 1)
|
|
image_encode = image_encode.astype(np.float32)
|
|
|
|
return image_encode
|
|
|
|
|
|
class BVTVertexMNN(HamburgerABC):
|
|
|
|
def __init__(self, mnn_path, *args, **kwargs):
|
|
from .common.mnn_adapt import MNNAdapter
|
|
super().__init__(*args, **kwargs)
|
|
self.input_shape = (1, 3, self.input_size[0], self.input_size[1])
|
|
self.session = MNNAdapter(mnn_path, self.input_shape)
|
|
|
|
def _run_session(self, data):
|
|
outputs = self.session.inference(data)
|
|
return outputs
|
|
|
|
def _postprocess(self, data):
|
|
assert data.shape[0] == 1
|
|
data = np.asarray(data).reshape(-1, 4, 2)
|
|
data[:, :, 0] *= self.input_size[1]
|
|
data[:, :, 1] *= self.input_size[0]
|
|
|
|
return data[0]
|
|
|
|
def _preprocess(self, image):
|
|
assert len(
|
|
image.shape) == 3, "The dimensions of the input image object do not match. The input supports a single " \
|
|
"image. "
|
|
image_resize = cv2.resize(image, self.input_size)
|
|
encode = encode_images(image_resize)
|
|
encode = encode.astype(np.float32)
|
|
input_tensor = np.expand_dims(encode, 0)
|
|
|
|
return input_tensor
|
|
|
|
|
|
class BVTVertexORT(HamburgerABC):
|
|
|
|
def __init__(self, onnx_path, *args, **kwargs):
|
|
import onnxruntime as ort
|
|
super().__init__(*args, **kwargs)
|
|
self.session = ort.InferenceSession(onnx_path, None)
|
|
self.input_config = self.session.get_inputs()[0]
|
|
self.output_config = self.session.get_outputs()[0]
|
|
self.input_size = self.input_config.shape[2:]
|
|
|
|
# @cost('Vertex')
|
|
def _run_session(self, data) -> np.ndarray:
|
|
result = self.session.run([self.output_config.name], {self.input_config.name: data})
|
|
|
|
return result[0]
|
|
|
|
def _postprocess(self, data) -> np.ndarray:
|
|
assert data.shape[0] == 1
|
|
data = np.asarray(data).reshape(-1, 4, 2)
|
|
data[:, :, 0] *= self.input_size[1]
|
|
data[:, :, 1] *= self.input_size[0]
|
|
|
|
return data[0]
|
|
|
|
def _preprocess(self, image) -> np.ndarray:
|
|
assert len(
|
|
image.shape) == 3, "The dimensions of the input image object do not match. The input supports a single " \
|
|
"image. "
|
|
image_resize = cv2.resize(image, self.input_size)
|
|
encode = encode_images(image_resize)
|
|
encode = encode.astype(np.float32)
|
|
input_tensor = np.expand_dims(encode, 0)
|
|
|
|
return input_tensor
|