前言

Hello大家好，今天給大家分享一下如何基于深度學(xué)習(xí)模型訓(xùn)練實(shí)現(xiàn)工件切割點(diǎn)位置預(yù)測(cè)，主要是通過(guò)對(duì)YOLOv8姿態(tài)評(píng)估模型在自定義的數(shù)據(jù)集上訓(xùn)練，生成一個(gè)工件切割分離點(diǎn)預(yù)測(cè)模型

01制作數(shù)據(jù)集

本人從網(wǎng)絡(luò)上隨便找到了個(gè)工業(yè)工件，然后寫(xiě)代碼合成了一些數(shù)據(jù)，總計(jì)數(shù)據(jù)有360張圖像、其中336張作為訓(xùn)練集、24張作為驗(yàn)證集。

其中YOLOv的數(shù)據(jù)格式如下：

解釋一下：

Class-index 表示對(duì)象類(lèi)型索引，從0開(kāi)始
后面的四個(gè)分別是對(duì)象的中心位置與寬高 xc、yc、width、height
Px1,py1表示第一個(gè)關(guān)鍵點(diǎn)坐標(biāo)、p1v表示師傅可見(jiàn)，默認(rèn)填2即可。

02模型訓(xùn)練

跟訓(xùn)練YOLOv8對(duì)象檢測(cè)模型類(lèi)似，直接運(yùn)行下面的命令行即可：

yolo train model=yolov8n-pose.pt data=mul_lines_dataset.yaml epochs=15 imgsz=640 batch=1

03模型導(dǎo)出預(yù)測(cè)

訓(xùn)練完成以后模型預(yù)測(cè)推理測(cè)試 使用下面的命令行：

yolo predict model=lines_pts_best.pt source=D:ird_testack12.png

導(dǎo)出模型為ONNX格式，使用下面命令行即可

yolo export model=lines_pts_best.pt format=onnx

04部署推理

基于ONNX格式模型，采用ONNXRUNTIME推理結(jié)果如下

ORT相關(guān)的推理演示代碼如下：

defort_keypoint_demo():

#initializetheonnxruntimesessionbyloadingmodelinCUDAsupport
model_dir="lines_pts_best.onnx"
session=onnxruntime.InferenceSession(model_dir,providers=['CUDAExecutionProvider'])

#就改這里, 把RTSP的地址配到這邊就好啦，然后直接運(yùn)行，其它任何地方都不準(zhǔn)改！
#切記把 onnx文件放到跟這個(gè)python文件同一個(gè)文件夾中！
frame=cv.imread("D:/bird_test/back1/lines_002.png")
bgr=format_yolov8(frame)
fh,fw,fc=frame.shape

start=time.time()
image=cv.dnn.blobFromImage(bgr,1/255.0,(640,640),swapRB=True,crop=False)

#onnxruntimeinference
ort_inputs={session.get_inputs()[0].name:image}
res=session.run(None,ort_inputs)[0]

#matrixtransposefrom1x8x8400=>8400x8
out_prob=np.squeeze(res,0).T

result_kypts,confidences,boxes=wrap_detection(bgr,out_prob)
for(kpts,confidence,box)inzip(result_kypts,confidences,boxes):
cv.rectangle(frame,box,(0,0,255),2)
cv.rectangle(frame,(box[0],box[1]-20),(box[0]+box[2],box[1]),(0,255,255),-1)
cv.putText(frame,("%.2f"%confidence),(box[0],box[1]-10),cv.FONT_HERSHEY_SIMPLEX,.5,(0,0,0))
cx=kpts[0]
cy=kpts[1]
cv.circle(frame,(int(cx),int(cy)),3,(255,0,255),4,8,0)

cv.imshow("Find Key PointDemo",frame)
cv.waitKey(0)
cv.destroyAllWindows()


if__name__=="__main__":
ort_keypoint_demo()

審核編輯：湯梓紅