onnx yolov8添加nms說(shuō)明文檔

官網(wǎng)下載源碼、模型和測(cè)試數(shù)據(jù)集：源代碼：/ultralytics/ultralytics模型：/ultralytics/assets/releases/download/v0.0.0/yolov8m.pt測(cè)試數(shù)據(jù)：/s/1q8t6WUe2r9NavY-4jYOokA?pwd=mf4i安裝依賴(lài)環(huán)境：cdultralyticspipinstall-rrequirements.txt-i/simplepythonsetup.pyinstall修改腳本（ultralytics/nn/tasks.py）defxywh2xyxy(x):

y=x.clone()

y[:,:,0]=x[:,:,0]-x[:,:,2]/2

#topleftx

y[:,:,1]=x[:,:,1]-x[:,:,3]/2

#toplefty

y[:,:,2]=x[:,:,0]+x[:,:,2]/2

#bottomrightx

y[:,:,3]=x[:,:,1]+x[:,:,3]/2

#bottomrighty

returnydefget_box(input_feature_np):

boxes_feat=input_feature_np[:,:,0:4]

boxes_xyxy=xywh2xyxy(boxes_feat)

returnboxes_xyxydefget_score(input_feature_np):

#confidence_feat=input_feature_np[:,:,4:5]

#classes_feat=input_feature_np[:,:,5:]

#score=confidence_feat*classes_feat

score=input_feature_np[:,:,4:]

returnscore

def_predict_once(self,x,profile=False,visualize=False):

"""

Performaforwardpassthroughthenetwork.

Args:

x(torch.Tensor):Theinputtensortothemodel.

profile(bool):

PrintthecomputationtimeofeachlayerifTrue,defaultstoFalse.

visualize(bool):SavethefeaturemapsofthemodelifTrue,defaultstoFalse

Returns:

(torch.Tensor):Thelastoutputofthemodel.

"""

y,dt=[],[]

#outputs

forminself.model:

ifm.f!=-1:

#ifnotfrompreviouslayer

x=y[m.f]ifisinstance(m.f,int)else[xifj==-1elsey[j]forjinm.f]

#fromearlierlayers

ifprofile:

self._profile_one_layer(m,x,dt)

x=m(x)

#run

y.append(xifm.iinself.saveelseNone)

#saveoutput

ifvisualize:

feature_visualization(x,m.type,m.i,save_dir=visualize)

xt=x.transpose(1,2)

bboxes=get_box(xt)

scores=get_score(xt)

returnbboxes,scores

returnx修改輸出（ultralytics/engine/exporter.py）name添加導(dǎo)出onnx腳本importonnxfromonnximportGraphProto,TensorProtofromonnx.helperimportmake_tensor_value_info,make_attribute,make_node,make_modelfromultralyticsimportYOLOdefappend_nms(graph,unused_node=[]):

ngraph=GraphProto()

=

ngraph.input.extend([iforiingraph.inputifnotinunused_node])

ngraph.initializer.extend([iforiingraph.initializerifnotinunused_node])

ngraph.value_info.extend([iforiingraph.value_infoifnotinunused_node])

ngraph.node.extend([iforiingraph.nodeifnotinunused_node])

output_info=[iforiingraph.output]

ngraph.value_info.extend(output_info)

print(graph.output)

nms=make_node(

'DLNonMaxSuppression',

inputs=['bboxes','scores'],

outputs=['num_detections','nmsed_boxes','nmsed_scores','nmsed_classes']

)

nms.attribute.append(make_attribute('backgroundLabelId',-1))

nms.attribute.append(make_attribute('iouThreshold',0.45))

nms.attribute.append(make_attribute('isNormalized',False))

nms.attribute.append(make_attribute('keepTopK',200))

nms.attribute.append(make_attribute('numClasses',80))#

nms.attribute.append(make_attribute('shareLocation',True))

nms.attribute.append(make_attribute('scoreThreshold',0.25))

nms.attribute.append(make_attribute('topK',1000))

ngraph.node.append(nms)

num_detection=make_tensor_value_info('num_detections',TensorProto.INT32,[1,1])

nmsed_box=make_tensor_value_info('nmsed_boxes',TensorProto.FLOAT,[1,200,4])

nmsed_score=make_tensor_value_info('nmsed_scores',TensorProto.FLOAT,[1,200,1])

nmsed_class=make_tensor_value_info('nmsed_classes',TensorProto.FLOAT,[1,200,1])

ngraph.output.extend([num_detection,nmsed_box,nmsed_score,nmsed_class])

returnngraphdefexport_onnx():

#Loadamodel

model=YOLO("yolov8m.pt")

#loadapretrainedmodel(recommendedfortraining)

f=model.export(format="onnx")

#exportthemodeltoONNXformat

print('model.exportsuccess')

#Checks

model_onnx=onnx.load(f)

#loadonnxmodel

onnx.checker.check_model(model_onnx)

#checkonnxmodel

#Simplify

try:

importonnxsim

model_onnx,check=onnxsim.simplify(model_onnx)

assertcheck,'assertcheckfailed'

onnx.save(model_onnx,f)

exceptExceptionase:

print(f'simplifierfailure:{e}')

returnmodel_onnxmodel=export_onnx()model_attrs=dict(

ir_version=