55 lines
1.4 KiB
Python
55 lines
1.4 KiB
Python
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)
|
||
|