基于YOLOv8實現自定義姿態評估模型訓練

前言

Hello大家好，今天給大家分享一下如何基于YOLOv8姿態評估模型，實現在自定義數據集上，完成自定義姿態評估模型的訓練與推理。

01tiger-pose數據集

YOLOv8官方提供了一個自定義tiger-pose數據集（老虎姿態評估），總計數據有263張圖像、其中210張作為訓練集、53張作為驗證集。

其中YOLOv8-pose的數據格式如下：

解釋一下：

Class-index 表示對象類型索引，從0開始
后面的四個分別是對象的中心位置與寬高 xc、yc、width、height
px1,py1表示第一個關鍵點坐標、p1v表示師傅可見，默認填2即可。
kpt_shape=12x2表示有12個關鍵點，每個關鍵點是x,y

02模型訓練

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

yolotrainmodel=yolov8n-pose.ptdata=tiger_pose_dataset.yamlepochs=100imgsz=640batch=1

03模型導出預測

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

yolo predict model=tiger_pose_best.pt source=D:/123.jpg

導出模型為ONNX格式，使用下面命令行即可

yolo export model=tiger_pose_best.pt format=onnx

04部署推理

基于ONNX格式模型，采用ONNXRUNTIME推理結果如下：

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

def ort_pose_demo():


    # initialize the onnxruntime session by loading model in CUDA support
    model_dir = "tiger_pose_best.onnx"
    session = onnxruntime.InferenceSession(model_dir, providers=['CUDAExecutionProvider'])


    # 就改這里, 把RTSP的地址配到這邊就好啦，然后直接運行，其它任何地方都不準改！
    # 切記把 yolov8-pose.onnx文件放到跟這個python文件同一個文件夾中！
    frame = cv.imread("D:/123.jpg")
    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))
        cv.circle(frame, (int(kpts[0]), int(kpts[1])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[2]), int(kpts[3])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[4]), int(kpts[5])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[6]), int(kpts[7])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[8]), int(kpts[9])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[10]), int(kpts[11])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[12]), int(kpts[13])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[14]), int(kpts[15])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[16]), int(kpts[17])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[18]), int(kpts[19])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[20]), int(kpts[21])), 3, (255, 0, 255), 4, 8, 0)
        cv.circle(frame, (int(kpts[22]), int(kpts[23])), 3, (255, 0, 255), 4, 8, 0)


    cv.imshow("Tiger Pose Demo - gloomyfish", frame)
    cv.waitKey(0)
    cv.destroyAllWindows()

審核編輯：湯梓紅