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()
-
C++
+關(guān)注
關(guān)注
21文章
2090瀏覽量
73410 -
pytorch
+關(guān)注
關(guān)注
2文章
795瀏覽量
13084 -
訓(xùn)練模型
+關(guān)注
關(guān)注
1文章
35瀏覽量
3794
原文標(biāo)題:TorchVision對(duì)象檢測(cè)RetinaNet推理演示
文章出處:【微信號(hào):CVSCHOOL,微信公眾號(hào):OpenCV學(xué)堂】歡迎添加關(guān)注!文章轉(zhuǎn)載請(qǐng)注明出處。
發(fā)布評(píng)論請(qǐng)先 登錄
相關(guān)推薦
評(píng)論