基于COCO的預(yù)訓(xùn)練模型mAP對(duì)應(yīng)關(guān)系

torchvision對(duì)象檢測(cè)介紹

Pytorch1.11版本以上支持Torchvision高版本支持以下對(duì)象檢測(cè)模型的遷移學(xué)習(xí)：

- Faster-RCNN
- Mask-RCNN
- FCOS
- RetinaNet
- SSD
- KeyPointsRCNN

其中基于COCO的預(yù)訓(xùn)練模型mAP對(duì)應(yīng)關(guān)系如下：

最近一段時(shí)間本人已經(jīng)全部親測(cè)，都可以轉(zhuǎn)換為ONNX格式模型，都可以支持ONNXRUNTIME框架的Python版本與C++版本推理，本文以RetinaNet為例，演示了從模型下載到導(dǎo)出ONNX格式，然后基于ONNXRUNTIME推理的整個(gè)流程。

RetinaNet轉(zhuǎn)ONNX

把模型轉(zhuǎn)換為ONNX格式，Pytorch是原生支持的，只需要把通過(guò)torch.onnx.export接口，填上相關(guān)的參數(shù)，然后直接運(yùn)行就可以生成ONNX模型文件。相關(guān)的轉(zhuǎn)換代碼如下：

model=tv.models.detection.retinanet_resnet50_fpn(pretrained=True)
dummy_input=torch.randn(1,3,1333,800)
model.eval()
model(dummy_input)
im=torch.zeros(1,3,1333,800).to("cpu")
torch.onnx.export(model,im,
"retinanet_resnet50_fpn.onnx",
verbose=False,
opset_version=11,
training=torch.onnx.TrainingMode.EVAL,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={'input':{0:'batch',2:'height',3:'width'}}
)

運(yùn)行時(shí)候控制臺(tái)會(huì)有一系列的警告輸出，但是絕對(duì)不影響模型轉(zhuǎn)換，影響不影響精度我還沒(méi)做個(gè)仔細(xì)的對(duì)比。 模型轉(zhuǎn)換之后，可以直接查看模型的輸入與輸出結(jié)構(gòu)，圖示如下：

RetinaNet的ONNX格式推理

基于Python版本的ONNXRUNTIME完成推理演示，這個(gè)跟我之前寫(xiě)過(guò)一篇文章Faster-RCNN的ONNX推理演示非常相似，大概是去年寫(xiě)的，鏈接在這里： 代碼很簡(jiǎn)單，只有三十幾行，Python就是方便使用，這里最需要注意的是輸入圖像的預(yù)處理必須是RGB格式，需要歸一化到0~1之間。對(duì)得到的三個(gè)輸出層分別解析，就可以獲取到坐標(biāo)（boxes里面包含的實(shí)際坐標(biāo)，無(wú)需轉(zhuǎn)換），推理部分的代碼如下：

importonnxruntimeasort
importcv2ascv
importnumpyasnp
importtorchvision


coco_names={'0':'background','1':'person','2':'bicycle','3':'car','4':'motorcycle','5':'airplane','6':'bus',
'7':'train','8':'truck','9':'boat','10':'trafficlight','11':'firehydrant','13':'stopsign',
'14':'parkingmeter','15':'bench','16':'bird','17':'cat','18':'dog','19':'horse','20':'sheep',
'21':'cow','22':'elephant','23':'bear','24':'zebra','25':'giraffe','27':'backpack',
'28':'umbrella','31':'handbag','32':'tie','33':'suitcase','34':'frisbee','35':'skis',
'36':'snowboard','37':'sportsball','38':'kite','39':'baseballbat','40':'baseballglove',
'41':'skateboard','42':'surfboard','43':'tennisracket','44':'bottle','46':'wineglass',
'47':'cup','48':'fork','49':'knife','50':'spoon','51':'bowl','52':'banana','53':'apple',
'54':'sandwich','55':'orange','56':'broccoli','57':'carrot','58':'hotdog','59':'pizza',
'60':'donut','61':'cake','62':'chair','63':'couch','64':'pottedplant','65':'bed',
'67':'diningtable','70':'toilet','72':'tv','73':'laptop','74':'mouse','75':'remote',
'76':'keyboard','77':'cellphone','78':'microwave','79':'oven','80':'toaster','81':'sink',
'82':'refrigerator','84':'book','85':'clock','86':'vase','87':'scissors','88':'teddybear',
'89':'hairdrier','90':'toothbrush'}

transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

sess_options=ort.SessionOptions()
src=cv.imread("D:/images/mmc.png")
cv.namedWindow("Retina-NetDetectionDemo",cv.WINDOW_AUTOSIZE)
image=cv.cvtColor(src,cv.COLOR_BGR2RGB)
blob=transform(image)
c,h,w=blob.shape
input_x=blob.view(1,c,h,w)
defto_numpy(tensor):
returntensor.detach().cpu().numpy()iftensor.requires_gradelsetensor.cpu().numpy()

#computeONNXRuntimeoutputprediction
ort_inputs={ort_session.get_inputs()[0].name:to_numpy(input_x)}
ort_outs=ort_session.run(None,ort_inputs)
#(N,4)dimensionalarraycontainingtheabsolutebounding-box
boxes=ort_outs[0]
scores=ort_outs[1]
labels=ort_outs[2]
print(boxes.shape,boxes.dtype,labels.shape,labels.dtype,scores.shape,scores.dtype)

index=0
forx1,y1,x2,y2inboxes:
ifscores[index]>0.65:
cv.rectangle(src,(np.int32(x1),np.int32(y1)),
(np.int32(x2),np.int32(y2)),(140,199,0),2,8,0)
label_id=labels[index]
label_txt=coco_names[str(label_id)]
cv.putText(src,label_txt,(np.int32(x1),np.int32(y1)),cv.FONT_HERSHEY_SIMPLEX,0.75,(0,0,255),1)
index+=1
cv.imshow("Retina-NetDetectionDemo",src)
cv.imwrite("D:/mmc_result.png",src)
cv.waitKey(0)
cv.destroyAllWindows()