引言
在計算機視覺的各種應用中,人體姿態檢測是一項極具挑戰性的任務,它能夠幫助我們理解人體各部位的空間位置。本文將詳細介紹如何使用 YOLOv8 和 Python 實現一個人體姿態檢測系統,涵蓋模型加載、圖像預處理、姿態預測到結果可視化的全流程實現。本文只做了行走,站立,跳三種姿態判斷,需要其他姿態可自行添加。
效果
開發環境
- Python 3.8+
- PyTorch 1.7+
- OpenCV 4.5+
- ultralytics YOLOv8
1. 模型加載與初始化
首先,我們需要一個 YOLOv8 的預訓練模型,該模型可以從 Ultralytics 官方網站下載。加載模型的主要目的是將其調整到評估模式,確保在推斷時模型的表現為最優。
import torchclass YOLOv8Pose:def __init__(self, model_path, device='cpu', conf=0.25, iou=0.7):self.model = self.load_model(model_path, device)# 其他初始化設置def load_model(self, model_path, device):ckpt = torch.load(model_path, map_location=device)model = ckpt['model'].to(device).eval()return model
2. 圖像預處理
圖像預處理是檢測流程中不可或缺的一部分,我們需要調整圖像的大小以適應模型的輸入要求,同時進行歸一化處理。
import cv2 as cv
from ultralytics.data.augment import LetterBoxdef preprocess(self, img_path):im = cv.imread(img_path)im = self.letterbox(im)im = im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGBim = np.ascontiguousarray(im)im = torch.from_numpy(im).to(self.device).float() / 255.0return im
3. 推斷與后處理
使用加載的模型進行前向推斷,并對輸出的檢測結果進行處理。
def infer(self, img):preds = self.model(img)# 使用非極大抑制處理預測結果return predsdef postprocess(self, prediction):# 調整預測框,解析關鍵點return results
4. 結果可視化
對檢測到的姿態進行可視化,包括繪制邊框、關鍵點和骨骼連接。
def draw_results(self, image, results):# 使用 OpenCV 繪制結果return image
5. 整合與測試
將上述所有步驟整合到一個流程中,對指定的圖片進行處理并展示結果。
if __name__ == "__main__":yolov8 = YOLOv8Pose(model_path='yolov8s-pose.pt')img_path = 'path_to_image.png'img = yolov8.preprocess(img_path)prediction = yolov8.infer(img)results = yolov8.postprocess(prediction)final_image = yolov8.draw_results(img, results)cv.imshow('Detection Results', final_image)cv.waitKey(0)
源碼已上傳到github,需要源碼請私信或著評論區留下郵箱。
希望這篇博客能夠幫助你理解并實踐 YOLOv8 在人體姿態檢測上的應用!