126 lines
6.3 KiB
Python
126 lines
6.3 KiB
Python
import numpy as np
|
|
from .yolov5_postprocess import postprocess_
|
|
from base.preprocess import preprocess_, preprocess_hw
|
|
from base.postprocess import postprocess_hw
|
|
from base.inference import inference_
|
|
|
|
class Yolov5Runner:
|
|
def __init__(self, model_path, model_id, input_shape, grid20_path, grid40_path, grid80_path, num_classes, conf_thres, iou_thres, top_k_num, anchors=None, label_mapping=None, vanish_point=0.0, filter_large_box_ratio=1.0, agnostic=0.0, **kwargs):
|
|
# --------------- STEP 1 load model----------------
|
|
from kneron_utils.model import load_model
|
|
self.model, self._inference_type = load_model(model_path, **kwargs)
|
|
|
|
# --------------- doing initialization-------------------
|
|
self.input_shape = input_shape
|
|
# NOTE: stride can be an arg; currently hardcoded [8, 16, 32]
|
|
max_s = 32
|
|
assert input_shape[0] % max_s == 0 and input_shape[1] % max_s == 0
|
|
self.model_id = str(model_id)
|
|
self.conf_thres = conf_thres
|
|
self.iou_thres = iou_thres
|
|
self.top_k_num = top_k_num
|
|
self.num_classes = num_classes
|
|
self.vanish_point = vanish_point
|
|
self.filter_large_box_ratio = filter_large_box_ratio
|
|
self.agnostic = True if agnostic==1.0 else False
|
|
self.grids = [np.load(grid80_path), np.load(grid40_path), np.load(grid20_path)]
|
|
self.anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] if anchors==None else anchors
|
|
|
|
# --------------- STEP 2 get all parameter as dict-----------------
|
|
self.init_config = locals()
|
|
self.init_config.update(kwargs)
|
|
self.init_config['_type'] = self._inference_type
|
|
|
|
def run(self, image, rectangle_list=None, **kwargs):
|
|
|
|
# rectangle_list: List of rectangle [x1, y1, w, h, score]. where xy1=top-left.
|
|
if rectangle_list is not None:
|
|
return self.run_rectangle_list(image, rectangle_list, **kwargs)
|
|
|
|
# -----------------preprocess-------------------
|
|
pre_config = {
|
|
'input_shape': self.input_shape,
|
|
'keep_ap': True,
|
|
'upscale': True,
|
|
}
|
|
pre_config.update(self.init_config)
|
|
pre_config.update(kwargs)
|
|
img_data, pre_info = preprocess_hw(image, preprocess_, **pre_config)#img_data:rgb
|
|
|
|
# -----------------do inference------------------
|
|
infer_config = {'model': self.model}
|
|
infer_config.update(self.init_config)
|
|
infer_config.update(pre_config)
|
|
outputs = inference_(img_data, **infer_config)
|
|
|
|
# -----------------post process------------------
|
|
post_config = {'model_id':self.model_id, 'conf_thres':self.conf_thres, 'iou_thres':self.iou_thres, 'top_k_num':self.top_k_num, 'grids':self.grids, 'num_classes':self.num_classes, 'vanish_point':self.vanish_point, 'filter_large_box_ratio':self.filter_large_box_ratio, 'agnostic':self.agnostic}
|
|
post_config.update(self.init_config)
|
|
post_config.update(pre_config)
|
|
post_config.update(pre_info)
|
|
post_config.update(kwargs)
|
|
post_config['anchors'] = self.anchors
|
|
|
|
outputs_onnx = postprocess_hw(outputs, postprocess_, **post_config)
|
|
# outputs_onnx : List of bboxes[x1, y1, w, h, score, class_id]. where xy1=top-left
|
|
label_mapping = post_config.get('label_mapping', None)
|
|
if label_mapping is not None:
|
|
for det in outputs_onnx:
|
|
det[-1] = label_mapping[int(det[-1]-1)]
|
|
return outputs_onnx
|
|
|
|
def run_rectangle_list(self, image, rectangle_list, **kwargs):
|
|
outputs_onnx_list = []
|
|
# 'rectangle_list': List of 'rectangle' [x1, y1, w, h, score]. where xy1=top-left.
|
|
# 'rectangle_list' = [[x1, y1, w, h, score], [x1, y1, w, h, score], ...].
|
|
for ii, rectangle in enumerate(rectangle_list):
|
|
# -----------------preprocess------------------
|
|
pre_config = {
|
|
'input_shape': self.input_shape,
|
|
'keep_ap': True,
|
|
'upscale': True,
|
|
'rectangle': rectangle,
|
|
}
|
|
pre_config.update(self.init_config)
|
|
pre_config.update(kwargs)
|
|
# 'img_data' is cropped from 'image' with 'rectangle' [x1, y1, w, h, score]
|
|
img_data, pre_info = preprocess_hw(image, preprocess_, **pre_config)
|
|
|
|
# -----------------do inference------------------
|
|
infer_config = {'model': self.model}
|
|
infer_config.update(self.init_config)
|
|
infer_config.update(pre_config)
|
|
outputs = inference_(img_data, **infer_config)
|
|
|
|
# -----------------post process------------------
|
|
post_config = {'model_id':self.model_id, 'conf_thres':self.conf_thres, 'iou_thres':self.iou_thres, 'top_k_num':self.top_k_num, 'grids':self.grids, 'num_classes':self.num_classes, 'vanish_point':self.vanish_point, 'filter_large_box_ratio':self.filter_large_box_ratio, 'agnostic':self.agnostic}
|
|
post_config.update(self.init_config)
|
|
post_config.update(pre_config)
|
|
post_config.update(pre_info)
|
|
post_config.update(kwargs)
|
|
post_config['anchors'] = self.anchors
|
|
outputs_onnx = postprocess_hw(outputs, postprocess_, **post_config)
|
|
# 'outputs_onnx' : List of bboxes[x1, y1, w, h, score, class_id]. where xy1=top-left
|
|
# 'outputs_onnx' = [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...].
|
|
if len(outputs_onnx)==0: continue
|
|
label_mapping = post_config.get('label_mapping', None)
|
|
if label_mapping is not None:
|
|
for det in outputs_onnx:
|
|
det[-1] = label_mapping[int(det[-1]-1)]
|
|
outputs_onnx_list.append(outputs_onnx)
|
|
# 'outputs_onnx_list' : List of outputs_onnx [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...].
|
|
# 'outputs_onnx_list' = [ [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...],
|
|
# [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...],
|
|
# ...
|
|
# [[x1, y1, w, h, score, class_id], [x1, y1, w, h, score, class_id], ...] ].
|
|
return outputs_onnx_list
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|