140 lines
5.2 KiB
Python
140 lines
5.2 KiB
Python
import cv2
|
||
import numpy as np
|
||
from ultralytics import YOLO
|
||
import time
|
||
import os
|
||
|
||
# ====================== 用户配置 ======================
|
||
MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/point/exp_pose2/weights/best.pt"
|
||
IMAGE_PATH = "/home/hx/yolo/output_masks/3.png"
|
||
SAVE_RESULT = True
|
||
SAVE_PATH = "/home/hx/yolo/test_image/result_with_angle2.jpg"
|
||
SHOW_IMAGE = True
|
||
|
||
# 关键点颜色：红、绿、蓝（对应 kpt 1, 2, 3）
|
||
KPT_COLORS = [
|
||
(0, 0, 255), # 红色 - kpt 1
|
||
(0, 255, 0), # 绿色 - kpt 2
|
||
(255, 0, 0), # 蓝色 - kpt 3
|
||
]
|
||
|
||
# ====================== 计算两向量夹角（0~180°）=======================
|
||
def calculate_angle_between_vectors(v1, v2):
|
||
"""
|
||
计算两个向量之间的夹角（单位：度），范围 0 ~ 180°
|
||
"""
|
||
v1 = np.array(v1)
|
||
v2 = np.array(v2)
|
||
|
||
if np.linalg.norm(v1) == 0 or np.linalg.norm(v2) == 0:
|
||
return 0.0
|
||
|
||
v1_u = v1 / np.linalg.norm(v1)
|
||
v2_u = v2 / np.linalg.norm(v2)
|
||
|
||
dot = np.clip(np.dot(v1_u, v2_u), -1.0, 1.0)
|
||
angle_rad = np.arccos(dot)
|
||
angle_deg = np.degrees(angle_rad)
|
||
return angle_deg
|
||
|
||
# ====================== 主程序 =======================
|
||
if __name__ == "__main__":
|
||
print("正在加载模型...")
|
||
try:
|
||
model = YOLO(MODEL_PATH)
|
||
print(f"✅ 模型加载完成: {MODEL_PATH}")
|
||
except Exception as e:
|
||
print(f"❌ 模型加载失败: {e}")
|
||
exit(1)
|
||
|
||
# 读取图像
|
||
img = cv2.imread(IMAGE_PATH)
|
||
if img is None:
|
||
raise FileNotFoundError(f"❌ 无法加载图像: {IMAGE_PATH}")
|
||
|
||
h, w = img.shape[:2]
|
||
print(f"✅ 图像加载成功: {IMAGE_PATH} ({w}x{h})")
|
||
|
||
# --- 推理 ---
|
||
results = model(img, conf=0.25, imgsz=1280)
|
||
result = results[0]
|
||
|
||
# 存储每个实例的关键点和置信度
|
||
instances = []
|
||
|
||
if result.keypoints is not None and result.boxes is not None:
|
||
kpts_list = result.keypoints.xy.cpu().numpy() # (N, K, 2)
|
||
confs = result.boxes.conf.cpu().numpy() # (N,)
|
||
bboxes = result.boxes.xyxy.cpu().numpy() # (N, 4)
|
||
|
||
# 组合：置信度 + 关键点 + 边框
|
||
for i in range(len(kpts_list)):
|
||
if len(kpts_list[i]) >= 2: # 至少有两个关键点
|
||
kpt1 = (int(kpts_list[i][0][0]), int(kpts_list[i][0][1]))
|
||
kpt2 = (int(kpts_list[i][1][0]), int(kpts_list[i][1][1]))
|
||
instances.append({
|
||
'conf': confs[i],
|
||
'kpt1': kpt1,
|
||
'kpt2': kpt2,
|
||
'bbox': bboxes[i]
|
||
})
|
||
|
||
# 按置信度降序排序
|
||
instances.sort(key=lambda x: x['conf'], reverse=True)
|
||
|
||
# 只保留置信度最高的两个实例
|
||
top_two = instances[:2]
|
||
|
||
if len(top_two) == 2:
|
||
inst1, inst2 = top_two
|
||
|
||
# 绘制第一个实例的关键点和连线
|
||
cv2.circle(img, inst1['kpt1'], 10, KPT_COLORS[0], -1)
|
||
cv2.circle(img, inst1['kpt2'], 10, KPT_COLORS[1], -1)
|
||
cv2.line(img, inst1['kpt1'], inst1['kpt2'], (255, 255, 0), 3) # 黄色线
|
||
cv2.putText(img, "1", (inst1['kpt1'][0]+10, inst1['kpt1'][1]-10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 1.0, KPT_COLORS[0], 2)
|
||
cv2.putText(img, "2", (inst1['kpt2'][0]+10, inst1['kpt2'][1]-10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 1.0, KPT_COLORS[1], 2)
|
||
|
||
# 绘制第二个实例的关键点和连线
|
||
cv2.circle(img, inst2['kpt1'], 10, KPT_COLORS[0], -1)
|
||
cv2.circle(img, inst2['kpt2'], 10, KPT_COLORS[1], -1)
|
||
cv2.line(img, inst2['kpt1'], inst2['kpt2'], (0, 255, 255), 3) # 青色线
|
||
cv2.putText(img, "1", (inst2['kpt1'][0]+10, inst2['kpt1'][1]-10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 1.0, KPT_COLORS[0], 2)
|
||
cv2.putText(img, "2", (inst2['kpt2'][0]+10, inst2['kpt2'][1]-10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 1.0, KPT_COLORS[1], 2)
|
||
|
||
# 构造两个向量：kpt2 - kpt1
|
||
vec1 = np.array([inst1['kpt2'][0] - inst1['kpt1'][0],
|
||
inst1['kpt2'][1] - inst1['kpt1'][1]])
|
||
|
||
vec2 = np.array([inst2['kpt2'][0] - inst2['kpt1'][0],
|
||
inst2['kpt2'][1] - inst2['kpt1'][1]])
|
||
|
||
# 计算夹角
|
||
angle = calculate_angle_between_vectors(vec1, vec2)
|
||
print(f"✅ 置信度最高的两个实例，其关键点1-2连线夹角: {angle:.1f}°")
|
||
|
||
# 在图像上显示角度
|
||
cv2.putText(img, f"Angle: {angle:.1f}°", (50, 50),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 3, lineType=cv2.LINE_AA)
|
||
|
||
else:
|
||
print(f"⚠️ 检测到 {len(top_two)} 个有效实例（需要至少2个）")
|
||
else:
|
||
print("⚠️ 未检测到关键点或边界框")
|
||
|
||
# --- 保存结果 ---
|
||
os.makedirs(os.path.dirname(SAVE_PATH), exist_ok=True)
|
||
if SAVE_RESULT:
|
||
cv2.imwrite(SAVE_PATH, img)
|
||
print(f"✅ 结果图像已保存至: {SAVE_PATH}")
|
||
|
||
# --- 显示图像 ---
|
||
if SHOW_IMAGE:
|
||
cv2.imshow("Keypoint Angle Between Two Best Instances", img)
|
||
print("👉 按任意键关闭窗口...")
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows() |