from ultralytics import YOLO
import cv2
import numpy as np
import os

# 1. 加载模型
model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb2/weights/best.pt')

# 2. 读取图像
img_path = r"/home/hx/yolo/output_masks/2.jpg"
img = cv2.imread(img_path)

if img is None:
    print(f"❌ 错误：无法读取图像！请检查路径：{img_path}")
    exit(1)

# 3. 预测（OBB 模式）
results = model(
    img,
    save=False,
    imgsz=640,
    conf=0.15,
    mode='obb'
)

# 4. 获取结果并绘制
result = results[0]
annotated_img = result.plot()

# 5. 保存结果
output_dir = "./inference_results"
os.makedirs(output_dir, exist_ok=True)
filename = os.path.basename(img_path)
save_path = os.path.join(output_dir, "detected_" + filename)
cv2.imwrite(save_path, annotated_img)
print(f"✅ 推理结果已保存至: {save_path}")

# 6. 提取旋转框并计算 **两个框之间的夹角**
boxes = result.obb
if boxes is None or len(boxes) == 0:
    print("❌ No objects detected.")
else:
    print(f"✅ Detected {len(boxes)} object(s):")
    directions = []  # 存储每个框的主方向（弧度），归一化到 [0, π)

    for i, box in enumerate(boxes):
        cls = int(box.cls.cpu().numpy()[0])
        conf = box.conf.cpu().numpy()[0]
        xywhr = box.xywhr.cpu().numpy()[0]  # [cx, cy, w, h, r]
        cx, cy, w, h, r_rad = xywhr

        # 确定主方向（长边方向）
        if w >= h:
            direction = r_rad  # 长边方向就是 r
        else:
            direction = r_rad + np.pi / 2  # 长边方向是 r + 90°

        # 归一化到 [0, π)
        direction = direction % np.pi

        directions.append(direction)
        angle_deg = np.degrees(direction)
        print(f"  Box {i+1}: Class: {cls}, Confidence: {conf:.3f}, 主方向: {angle_deg:.2f}°")

    # ✅ 计算任意两个框之间的夹角（最小夹角，0° ~ 90°）
    if len(directions) >= 2:
        print("\n🔍 计算两个旋转框之间的夹角（主方向夹角）：")
        for i in range(len(directions)):
            for j in range(i + 1, len(directions)):
                dir1 = directions[i]
                dir2 = directions[j]

                # 计算方向差（取最小夹角，考虑周期性）
                diff = abs(dir1 - dir2)
                diff = min(diff, np.pi - diff)  # 最小夹角（0 ~ π/2）
                diff_deg = np.degrees(diff)

                print(f"  Box {i+1} 与 Box {j+1} 之间的夹角: {diff_deg:.2f}°")
    else:
        print("⚠️  检测到少于两个目标，无法计算夹角。")

# 7. 显示图像
cv2.imshow("YOLO OBB Prediction", annotated_img)
cv2.waitKey(0)
cv2.destroyAllWindows()