YOLOV8 POSE姿态检测对图片绘制矩形和和关节点序号

西柚与蓝莓2024-08-21 21:40

代码如下

复制代码

import cv2
import torch

from ultralytics import YOLO

# Load a model
# model = YOLO("yolov8n-pose.yaml")  # build a new model from YAML
model = YOLO("yolov8n-pose.pt")  # load a pretrained model (recommended for training)
# model = YOLO("yolov8n-pose.yaml").load("yolov8n-pose.pt")  # build from YAML and transfer weights

# Train the model
import cv2

cv2.namedWindow('win', cv2.WINDOW_NORMAL)
# cv2.resizeWindow('win', 400, 600)
img_path=r"bus.jpg"
img=cv2.imread(img_path)
results = model(img_path)
for ret in results[0]:
    xmin,ymin,xmax,ymax,conf,cls=ret.boxes.data[0]
    cv2.rectangle(img,(int(xmin),int(ymin)),(int(xmax),int(ymax)),color=(255,0,0),thickness=3,lineType=cv2.LINE_AA)
    keypoints=ret.keypoints.data[0]
    for idx,pt in enumerate(keypoints):
        x,y=int(pt[0]),int(pt[1])
        cv2.circle(img,center=(x,y),radius=2,color=(0,0,255),thickness=2)
        cv2.putText(img,str(idx),org=(x-8,y),fontFace=cv2.FONT_ITALIC,fontScale=0.5,color=(0,0,255),thickness=1)
cv2.imshow("win",img)
cv2.waitKey(0)
cv2.destroyAllWindows()
# print(results)

效果如下

每个可以看得见的关节都标注了序号