前言

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

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

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

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

解釋一下：

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

02模型訓練

跟訓練YOLOv8對象檢測模型類似，直接運行下面的命令行即可：

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

03模型導出預測

訓練完成以后模型預測推理測試 使用下面的命令行：

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

導出模型為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的地址配到這邊就好啦，然后直接運行，其它任何地方都不準改！
#切記把 onnx文件放到跟這個python文件同一個文件夾中！
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()

審核編輯：湯梓紅