onnx添加nms

最新推荐文章于 2024-10-17 19:20:37 发布

wuqingshan2010

最新推荐文章于 2024-10-17 19:20:37 发布

阅读量3.9k

点赞数 3

CC 4.0 BY-SA版权

分类专栏： onnx 目标检测

本文链接：https://blog.csdn.net/wuqingshan2010/article/details/109471454

import numpy as np

import onnx
from onnx import defs, checker, helper, numpy_helper, mapping
from onnx import ModelProto, GraphProto, NodeProto, AttributeProto, TensorProto, OperatorProto, OperatorSetIdProto
from onnx.helper import make_tensor, make_tensor_value_info, make_attribute, make_model, make_node

dynamic_batch = False

def append_nms(graph, unused_node=[]):
    ngraph = GraphProto()
    ngraph.name = graph.name

    ngraph.input.extend([i for i in graph.input if i.name not in unused_node])
    ngraph.initializer.extend([i for i in graph.initializer if i.name not in unused_node])
    ngraph.value_info.extend([i for i in graph.value_info if i.name not in unused_node])
    ngraph.node.extend([i for i in graph.node if i.name not in unused_node])

    output_info = [i for i in graph.output]
    ngraph.value_info.extend(output_info)
    # print(graph.output)

    score_node = 'scores'
    bbox_node = 'boxes'

    # ngraph.value_info.append(make_tensor_value_info(score_n