import cv2
import numpy as np
from rtmlib import Body, draw_skeleton

def detect_body_keypoints_image(image_path, output_path=None):
    # 初始化模型
    device = 'cuda' if cv2.cuda.getCudaEnabledDeviceCount() > 0 else 'cpu'
    backend = 'onnxruntime'
    openpose_skeleton = True  # True为OpenPose风格，False为MMPose风格
    
    # 创建Body实例，使用平衡模式
    body = Body(
        mode='balanced',
        to_openpose=openpose_skeleton,
        backend=backend,
        device=device
    )
    
    # 读取图片
    image = cv2.imread(image_path)
    if image is None:
        print(f"无法读取图片: {image_path}")
        return
    
    # 检测身体关键点
    keypoints, scores = body(image)

    print(keypoints)
    
    # 绘制关键点
    result_image = draw_skeleton(
        image.copy(),
        keypoints,
        scores,
        openpose_skeleton=openpose_skeleton,
        kpt_thr=0.43
    )
    
    # 保存或显示结果
    if output_path:
        cv2.imwrite(output_path, result_image)
        print(f"已保存结果到: {output_path}")
    
    # 显示结果
    # cv2.imshow('Body Keypoints Detection', result_image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

if __name__ == "__main__":
    # 替换为您的图片路径
    image_path = "example.png"
    output_path = "body_result.jpg"
    detect_body_keypoints_image(image_path, output_path)