前言

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

01制作數據集

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

其中YOLOv的數據格式如下：

解釋一下：

Class-index 表示對象類型索引，從0開始
后面的四個分別是對象的中心位置與寬高 xc、yc、width、height
Px1,py1表示第一個關鍵點坐標、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推理結果如下

ORT相關的推理演示代碼如下：

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()

審核編輯：湯梓紅