未验证 提交 b41194ea 编写于 作者: K Kevin_Xiong 提交者: GitHub

add flag skipping postprocess to support edgeboard hardware (#6719)

* add flag skipping postprocess to support edgeboard hardware

* add flag skipping postprocess to support edgeboard hardware

* add flag skipping postprocess to support edgeboard hardware

* add comment for the flag export_eb
上级 e567dfe3
...@@ -131,12 +131,16 @@ def _dump_infer_config(config, path, image_shape, model): ...@@ -131,12 +131,16 @@ def _dump_infer_config(config, path, image_shape, model):
'use_dynamic_shape': use_dynamic_shape 'use_dynamic_shape': use_dynamic_shape
}) })
export_onnx = config.get('export_onnx', False) export_onnx = config.get('export_onnx', False)
export_eb = config.get('export_eb', False)
infer_arch = config['architecture'] infer_arch = config['architecture']
if 'RCNN' in infer_arch and export_onnx: if 'RCNN' in infer_arch and export_onnx:
logger.warning( logger.warning(
"Exporting RCNN model to ONNX only support batch_size = 1") "Exporting RCNN model to ONNX only support batch_size = 1")
infer_cfg['export_onnx'] = True infer_cfg['export_onnx'] = True
infer_cfg['export_eb'] = export_eb
if infer_arch in MOT_ARCH: if infer_arch in MOT_ARCH:
if infer_arch == 'DeepSORT': if infer_arch == 'DeepSORT':
......
...@@ -34,16 +34,17 @@ __all__ = [ ...@@ -34,16 +34,17 @@ __all__ = [
@register @register
class BBoxPostProcess(object): class BBoxPostProcess(object):
__shared__ = ['num_classes', 'export_onnx'] __shared__ = ['num_classes', 'export_onnx', 'export_eb']
__inject__ = ['decode', 'nms'] __inject__ = ['decode', 'nms']
def __init__(self, num_classes=80, decode=None, nms=None, def __init__(self, num_classes=80, decode=None, nms=None,
export_onnx=False): export_onnx=False, export_eb=False):
super(BBoxPostProcess, self).__init__() super(BBoxPostProcess, self).__init__()
self.num_classes = num_classes self.num_classes = num_classes
self.decode = decode self.decode = decode
self.nms = nms self.nms = nms
self.export_onnx = export_onnx self.export_onnx = export_onnx
self.export_eb = export_eb
def __call__(self, head_out, rois, im_shape, scale_factor): def __call__(self, head_out, rois, im_shape, scale_factor):
""" """
...@@ -100,6 +101,10 @@ class BBoxPostProcess(object): ...@@ -100,6 +101,10 @@ class BBoxPostProcess(object):
pred_result (Tensor): The final prediction results with shape [N, 6] pred_result (Tensor): The final prediction results with shape [N, 6]
including labels, scores and bboxes. including labels, scores and bboxes.
""" """
if self.export_eb:
# enable rcnn models for edgeboard hw to skip the following postprocess.
return bboxes, bboxes, bbox_num
if not self.export_onnx: if not self.export_onnx:
bboxes_list = [] bboxes_list = []
bbox_num_list = [] bbox_num_list = []
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册