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_obb/weights/best.pt')

# 2. 读取图像
img_path = r"/output_masks/3.png"
img = cv2.imread(img_path)

# ✅ 检查图像是否加载成功
if img is None:
    print(f"❌ 错误：无法读取图像！请检查路径：{img_path}")
    print("💡 提示：可能是文件不存在、路径错误或图像损坏")
    exit(1)

# 3. 预测（使用 OBB 模式！）
results = model(
    img,
    save=False,
    imgsz=1280,      # 必须和训练时一致
    conf=0.25,
    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):")
    angles = []  # 存储每个框的旋转角度（角度制）

    for i, box in enumerate(boxes):
        cls = int(box.cls.cpu().numpy()[0])
        conf = box.conf.cpu().numpy()[0]

        # ✅ 正确方式：使用 .xywhr 获取旋转角度（新版本 API）
        xywhr = box.xywhr.cpu().numpy()[0]  # [x_center, y_center, width, height, rotation]
        angle_rad = xywhr[4]  # 第5个值是旋转角度（弧度）
        angle_deg = np.degrees(angle_rad)  # 转为角度

        angles.append(angle_deg)
        print(f"  Box {i+1}: Class: {cls}, Confidence: {conf:.3f}, Angle: {angle_deg:.2f}°")

    # ✅ 计算任意两个框之间的最小夹角差
    if len(angles) >= 2:
        print("\n🔍 计算旋转框之间的夹角差：")
        for i in range(len(angles)):
            for j in range(i + 1, len(angles)):
                diff = abs(angles[i] - angles[j])
                min_angle_diff = min(diff, 180 - diff)  # 取最小夹角（0~180°内）
                print(f"  Box {i+1} 与 Box {j+1} 的最小夹角差: {min_angle_diff:.2f}°")
    else:
        print("⚠️  检测到少于两个目标，无法计算夹角。")

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