點選上方 藍字 關註我們
微信公眾號： OpenCV學堂
關註獲取更多電腦視覺與深度學習知識
前言
Hello大家好，今天給大家分享一下如何基於深度學習模型訓練實作圓檢測與圓心位置預測，主要是透過對YOLOv8姿態評估模型在自訂的數據集上訓練，生成一個自訂的圓檢測與圓心定位預測模型
01
制作數據集
本人從網路上隨便找到了個工業工件，然後寫程式碼合成了一些數據，總計數據有360張影像、其中336張作為訓練集、24張作為驗證集。
其中YOLOv的數據格式如下：
解釋一下：
class-index 表示物件型別索引，從0開始後面的四個分別是物件的中心位置與寬高 xc、yc、width、heightPx1,py1表示第一個關鍵點座標、p1v表示師傅可見，預設填2即可。
02
模型訓練
跟訓練YOLOv8物件檢測模型類似，直接執行下面的命令列即可：
yolo train model=yolov8n-pose.pt data=circle_dataset.yaml epochs=15 imgsz=640 batch=1
03
模型匯出預測
訓練完成以後模型預測推理測試 使用下面的命令列：
yolo predict model=D:\python\my_yolov8_train_demo\runs\pose\train3\weights\best.pt source=D:\bird_test\back1\2.png
匯出模型為ONNX格式，使用下面命令列即可
yolo export model=D:\python\my_yolov8_train_demo\runs\pose\train3\weights\best.pt format=onnx
04
部署推理
基於ONNX格式模型，采用ONNXRUNTIME推理結果如下：
ORT相關的推理演示程式碼如下：
defort_circle_demo():
# initialize the onnxruntime session by loading model in CUDA support
model_dir = "D:/python/my_yolov8_train_demo/circle_detect.onnx"
session = onnxruntime.InferenceSession(model_dir, providers=['CUDAExecutionProvider'])
# 就改這裏, 把RTSP的地址配到這邊就好啦，然後直接執行，其它任何地方都不準改！
# 切記把 onnx檔放到跟這個python檔同一個資料夾中！
frame = cv.imread("D:/bird_test/back1/3.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)
# onnxruntime inference
ort_inputs = {session.get_inputs()[0].name: image}
res = session.run(None, ort_inputs)[0]
# matrix transpose from 1x8x8400 => 8400x8
out_prob = np.squeeze(res, 0).T
result_kypts, confidences, boxes = wrap_detection(bgr, out_prob)
for (kpts, confidence, box) in zip(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("Circle Detection Demo", frame)
cv.waitKey(0)
cv.destroyAllWindows()
if __name__ == "__main__":
ort_circle_demo()
掃碼學習YOLOv8視訊課程
推薦閱讀
