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.
185 lines
6.2 KiB
Python
185 lines
6.2 KiB
Python
import numpy as np
|
|
import cv2
|
|
import copy
|
|
from .base.base import HamburgerABC
|
|
|
|
|
|
def xywh2xyxy(boxes):
|
|
xywh = copy.deepcopy(boxes)
|
|
xywh[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
|
|
xywh[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
|
|
xywh[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
|
|
xywh[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
|
|
return xywh
|
|
|
|
|
|
def nms(boxes, iou_thresh): # nms
|
|
index = np.argsort(boxes[:, 4])[::-1]
|
|
keep = []
|
|
while index.size > 0:
|
|
i = index[0]
|
|
keep.append(i)
|
|
x1 = np.maximum(boxes[i, 0], boxes[index[1:], 0])
|
|
y1 = np.maximum(boxes[i, 1], boxes[index[1:], 1])
|
|
x2 = np.minimum(boxes[i, 2], boxes[index[1:], 2])
|
|
y2 = np.minimum(boxes[i, 3], boxes[index[1:], 3])
|
|
|
|
w = np.maximum(0, x2 - x1)
|
|
h = np.maximum(0, y2 - y1)
|
|
|
|
inter_area = w * h
|
|
union_area = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1]) + (
|
|
boxes[index[1:], 2] - boxes[index[1:], 0]) * (boxes[index[1:], 3] - boxes[index[1:], 1])
|
|
iou = inter_area / (union_area - inter_area)
|
|
idx = np.where(iou <= iou_thresh)[0]
|
|
index = index[idx + 1]
|
|
return keep
|
|
|
|
|
|
def restore_box(boxes, r, left, top):
|
|
boxes[:, [0, 2, 5, 7, 9, 11]] -= left
|
|
boxes[:, [1, 3, 6, 8, 10, 12]] -= top
|
|
|
|
boxes[:, [0, 2, 5, 7, 9, 11]] /= r
|
|
boxes[:, [1, 3, 6, 8, 10, 12]] /= r
|
|
return boxes
|
|
|
|
|
|
def detect_pre_precessing(img, img_size):
|
|
img, r, left, top = letter_box(img, img_size)
|
|
img = img[:, :, ::-1].transpose(2, 0, 1).copy().astype(np.float32)
|
|
img = img / 255
|
|
img = img.reshape(1, *img.shape)
|
|
return img, r, left, top
|
|
|
|
|
|
def post_precessing(dets, r, left, top, conf_thresh=0.25, iou_thresh=0.5):
|
|
choice = dets[:, :, 4] > conf_thresh
|
|
dets = dets[choice]
|
|
dets[:, 13:15] *= dets[:, 4:5]
|
|
box = dets[:, :4]
|
|
boxes = xywh2xyxy(box)
|
|
score = np.max(dets[:, 13:15], axis=-1, keepdims=True)
|
|
index = np.argmax(dets[:, 13:15], axis=-1).reshape(-1, 1)
|
|
output = np.concatenate((boxes, score, dets[:, 5:13], index), axis=1)
|
|
reserve_ = nms(output, iou_thresh)
|
|
output = output[reserve_]
|
|
output = restore_box(output, r, left, top)
|
|
return output
|
|
|
|
|
|
def letter_box(img, size=(640, 640)):
|
|
h, w, c = img.shape
|
|
r = min(size[0] / h, size[1] / w)
|
|
new_h, new_w = int(h * r), int(w * r)
|
|
top = int((size[0] - new_h) / 2)
|
|
left = int((size[1] - new_w) / 2)
|
|
|
|
bottom = size[0] - new_h - top
|
|
right = size[1] - new_w - left
|
|
img_resize = cv2.resize(img, (new_w, new_h))
|
|
img = cv2.copyMakeBorder(img_resize, top, bottom, left, right, borderType=cv2.BORDER_CONSTANT,
|
|
value=(0, 0, 0))
|
|
return img, r, left, top
|
|
|
|
|
|
class MultiTaskDetectorMNN(HamburgerABC):
|
|
|
|
def __init__(self, mnn_path, box_threshold: float = 0.5, nms_threshold: float = 0.6, *args, **kwargs):
|
|
from hyperlpr3.common.mnn_adapt import MNNAdapter
|
|
super().__init__(*args, **kwargs)
|
|
assert self.input_size[0] == self.input_size[1]
|
|
self.box_threshold = box_threshold
|
|
self.nms_threshold = nms_threshold
|
|
self.input_shape = (1, 3, self.input_size[0], self.input_size[1])
|
|
if self.input_size[0] == 320:
|
|
self.tensor_shape = [(1, 6300, 15)]
|
|
elif self.input_size[0] == 640:
|
|
self.tensor_shape = [(1, 25200, 15)]
|
|
self.session = MNNAdapter(mnn_path, self.input_shape, outputs_name=['output', ],
|
|
outputs_shape=self.tensor_shape)
|
|
|
|
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]))
|
|
result = np.asarray(result)
|
|
|
|
return result[0]
|
|
|
|
def _postprocess(self, data):
|
|
r, left, top = self.tmp_pack
|
|
|
|
return post_precessing(data, r, left, top)
|
|
|
|
def _preprocess(self, image):
|
|
img, r, left, top = detect_pre_precessing(image, self.input_size)
|
|
self.tmp_pack = r, left, top
|
|
|
|
return img
|
|
|
|
|
|
class MultiTaskDetectorDNN(HamburgerABC):
|
|
|
|
def __init__(self, onnx_path, box_threshold: float = 0.5, nms_threshold: float = 0.6, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.box_threshold = box_threshold
|
|
self.nms_threshold = nms_threshold
|
|
self.session = cv2.dnn.readNetFromONNX(onnx_path)
|
|
self.input_shape = (1, 3, self.input_size[0], self.input_size[1])
|
|
self.tensor_shape = [(1, 6300, 15)]
|
|
|
|
def _run_session(self, data):
|
|
self.session.setInput(data)
|
|
outputs = self.session.forward()
|
|
|
|
return outputs
|
|
|
|
def _postprocess(self, data):
|
|
|
|
r, left, top = self.tmp_pack
|
|
|
|
return post_precessing(data, r, left, top)
|
|
|
|
def _preprocess(self, image):
|
|
img, r, left, top = detect_pre_precessing(image, self.input_size)
|
|
self.tmp_pack = r, left, top
|
|
|
|
return img
|
|
|
|
|
|
class MultiTaskDetectorORT(HamburgerABC):
|
|
|
|
def __init__(self, onnx_path, box_threshold: float = 0.5, nms_threshold: float = 0.6, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
import onnxruntime as ort
|
|
self.box_threshold = box_threshold
|
|
self.nms_threshold = nms_threshold
|
|
self.session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
|
|
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
|
|
|
|
def _run_session(self, data):
|
|
result = self.session.run([self.outputs_option[0].name], {self.input_name: data})[0]
|
|
|
|
return result
|
|
|
|
def _postprocess(self, data):
|
|
r, left, top = self.tmp_pack
|
|
return post_precessing(data, r, left, top)
|
|
|
|
def _preprocess(self, image):
|
|
img, r, left, top = detect_pre_precessing(image, self.input_size)
|
|
self.tmp_pack = r, left, top
|
|
|
|
return img
|