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.
201 lines
7.2 KiB
Python
201 lines
7.2 KiB
Python
from hyperlpr3.common.tools_process import *
|
|
from .base.base import HamburgerABC
|
|
|
|
ANCHORS_MAP = {
|
|
320: [[9.38281, 3.08398], [15.53125, 4.93750], [19.98438, 7.78906],
|
|
[31.10938, 10.35156], [45.21875, 14.14844], [32.34375, 21.04688],
|
|
[65.62500, 19.57812], [76.12500, 46.12500], [253.25000, 137.50000]],
|
|
640: [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],
|
|
[59, 119], [116, 90], [156, 198], [373, 326]]
|
|
}
|
|
|
|
|
|
def image_to_input_tensor(image):
|
|
data = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
data = data.transpose(2, 0, 1) / 255.0
|
|
data = np.expand_dims(data, 0)
|
|
data = data.astype(np.float32)
|
|
|
|
return data
|
|
|
|
|
|
class Y5rkDetectorMNN(HamburgerABC):
|
|
def __init__(self, mnn_path, box_threshold: float = 0.5, nms_threshold: float = 0.6, *args, **kwargs):
|
|
from .common.mnn_adapt import MNNAdapter
|
|
super().__init__(*args, **kwargs)
|
|
self.box_threshold = box_threshold
|
|
self.nms_threshold = nms_threshold
|
|
self.input_shape = (1, 3, self.input_size[0], self.input_size[1])
|
|
self.tensor_shape = ((1, 18, 40, 40), (1, 18, 20, 20), (1, 18, 10, 10))
|
|
self.session = MNNAdapter(mnn_path, self.input_shape, outputs_name=['output', '335', '336'],
|
|
outputs_shape=self.tensor_shape)
|
|
assert self.input_size[0] == self.input_size[1]
|
|
self.anchors = ANCHORS_MAP[self.input_size[0]]
|
|
|
|
def _run_session(self, data):
|
|
outputs = self.session.inference(data)
|
|
result = list()
|
|
for idx, output in enumerate(outputs):
|
|
result.append(output.reshape(self.tensor_shape[idx]))
|
|
|
|
return result
|
|
|
|
def _postprocess(self, data):
|
|
ratio, (dw, dh) = self.temp_pack
|
|
input0_data = data[0]
|
|
input1_data = data[1]
|
|
input2_data = data[2]
|
|
input0_data = input0_data.reshape([3, -1] + list(input0_data.shape[-2:]))
|
|
input1_data = input1_data.reshape([3, -1] + list(input1_data.shape[-2:]))
|
|
input2_data = input2_data.reshape([3, -1] + list(input2_data.shape[-2:]))
|
|
input_data = list()
|
|
input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))
|
|
input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))
|
|
input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))
|
|
|
|
boxes, classes, scores = self.decode_outputs(input_data, self.input_size)
|
|
|
|
boxes = restore_bound_box(boxes, ratio, (dw, dh))
|
|
|
|
return boxes, classes, scores
|
|
|
|
def _preprocess(self, image):
|
|
h, w, _ = image.shape
|
|
resize_img, ratio, (dw, dh) = letterbox(image, new_shape=(self.input_size[1], self.input_size[0]))
|
|
data = image_to_input_tensor(resize_img)
|
|
self.temp_pack = ratio, (dw, dh)
|
|
|
|
return data
|
|
|
|
def decode_outputs(self, input_data, size):
|
|
masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
|
|
anchors = self.anchors
|
|
|
|
boxes, classes, scores = [], [], []
|
|
for input, mask in zip(input_data, masks):
|
|
b, c, s = process(input, mask, anchors, size)
|
|
b, c, s = filter_boxes(b, c, s, self.box_threshold, self.nms_threshold)
|
|
boxes.append(b)
|
|
classes.append(c)
|
|
scores.append(s)
|
|
|
|
boxes = np.concatenate(boxes)
|
|
boxes = xywh2xyxy(boxes)
|
|
classes = np.concatenate(classes)
|
|
scores = np.concatenate(scores)
|
|
|
|
nboxes, nclasses, nscores = [], [], []
|
|
for c in set(classes):
|
|
inds = np.where(classes == c)
|
|
b = boxes[inds]
|
|
c = classes[inds]
|
|
s = scores[inds]
|
|
|
|
keep = nms_boxes(b, s, self.nms_threshold)
|
|
|
|
nboxes.append(b[keep])
|
|
nclasses.append(c[keep])
|
|
nscores.append(s[keep])
|
|
|
|
if not nclasses and not nscores:
|
|
return None, None, None
|
|
|
|
boxes = np.concatenate(nboxes)
|
|
classes = np.concatenate(nclasses)
|
|
scores = np.concatenate(nscores)
|
|
|
|
return boxes, classes, scores
|
|
|
|
|
|
class Y5rkDetectorORT(HamburgerABC):
|
|
|
|
def __init__(self, onnx_path, box_threshold: float = 0.5, nms_threshold: float = 0.6, *args, **kwargs):
|
|
import onnxruntime as ort
|
|
super().__init__(*args, **kwargs)
|
|
self.box_threshold = box_threshold
|
|
self.nms_threshold = nms_threshold
|
|
self.session = ort.InferenceSession(onnx_path, None)
|
|
self.inputs_option = self.session.get_inputs()
|
|
self.outputs_option = self.session.get_outputs()
|
|
input_option = self.inputs_option[0]
|
|
input_size_ = tuple(input_option.shape[2:])
|
|
self.input_size = tuple(self.input_size)
|
|
if not self.input_size:
|
|
self.input_size = input_size_
|
|
assert self.input_size == input_size_, 'The dimensions of the input do not match the model expectations.'
|
|
assert self.input_size[0] == self.input_size[1]
|
|
self.input_name = input_option.name
|
|
self.anchors = ANCHORS_MAP[self.input_size[0]]
|
|
|
|
def decode_outputs(self, input_data, size):
|
|
masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
|
|
anchors = self.anchors
|
|
|
|
boxes, classes, scores = [], [], []
|
|
for input, mask in zip(input_data, masks):
|
|
b, c, s = process(input, mask, anchors, size)
|
|
b, c, s = filter_boxes(b, c, s, self.box_threshold, self.nms_threshold)
|
|
boxes.append(b)
|
|
classes.append(c)
|
|
scores.append(s)
|
|
|
|
boxes = np.concatenate(boxes)
|
|
boxes = xywh2xyxy(boxes)
|
|
classes = np.concatenate(classes)
|
|
scores = np.concatenate(scores)
|
|
|
|
nboxes, nclasses, nscores = [], [], []
|
|
for c in set(classes):
|
|
inds = np.where(classes == c)
|
|
b = boxes[inds]
|
|
c = classes[inds]
|
|
s = scores[inds]
|
|
|
|
keep = nms_boxes(b, s, self.nms_threshold)
|
|
|
|
nboxes.append(b[keep])
|
|
nclasses.append(c[keep])
|
|
nscores.append(s[keep])
|
|
|
|
if not nclasses and not nscores:
|
|
return None, None, None
|
|
|
|
boxes = np.concatenate(nboxes)
|
|
classes = np.concatenate(nclasses)
|
|
scores = np.concatenate(nscores)
|
|
|
|
return boxes, classes, scores
|
|
|
|
@cost("Detect")
|
|
def _run_session(self, data):
|
|
outputs = self.session.run([], {"images": data})
|
|
|
|
return outputs
|
|
|
|
def _postprocess(self, data):
|
|
ratio, (dw, dh) = self.temp_pack
|
|
input0_data = data[0]
|
|
input1_data = data[1]
|
|
input2_data = data[2]
|
|
input0_data = input0_data.reshape([3, -1] + list(input0_data.shape[-2:]))
|
|
input1_data = input1_data.reshape([3, -1] + list(input1_data.shape[-2:]))
|
|
input2_data = input2_data.reshape([3, -1] + list(input2_data.shape[-2:]))
|
|
input_data = list()
|
|
input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))
|
|
input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))
|
|
input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))
|
|
|
|
boxes, classes, scores = self.decode_outputs(input_data, self.input_size)
|
|
|
|
boxes = restore_bound_box(boxes, ratio, (dw, dh))
|
|
|
|
return boxes, classes, scores
|
|
|
|
def _preprocess(self, image):
|
|
h, w, _ = image.shape
|
|
resize_img, ratio, (dw, dh) = letterbox(image, new_shape=(self.input_size[1], self.input_size[0]))
|
|
data = image_to_input_tensor(resize_img)
|
|
self.temp_pack = ratio, (dw, dh)
|
|
|
|
return data
|