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

def get_best_obb_angle(image_path, weight_path, return_degree=False):
    """
    输入：
        image_path: 图像路径
        weight_path: YOLO权重路径
        return_degree: 是否返回角度单位为度，默认 False（返回弧度）
    输出：
        置信度最高目标的旋转角
        如果未检测到目标返回 None
    """
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        print(f"❌ 无法读取图像：{image_path}")
        return None

    # 加载模型并预测
    model = YOLO(weight_path)
    results = model(img, save=False, imgsz=640, conf=0.15, mode='obb')
    result = results[0]

    boxes = result.obb
    if not boxes:
        print("⚠️ 未检测到目标。")
        return None

    # 取置信度最高框的旋转角
    best_box = max(boxes, key=lambda x: x.conf.cpu().numpy()[0])
    r = best_box.xywhr.cpu().numpy()[0][4]  # 弧度

    if return_degree:
        return np.degrees(r)
    else:
        return r


def save_obb_visual(image_path, weight_path, save_path):
    """
    输入：
        image_path: 图像路径
        weight_path: YOLO权重路径
        save_path: 保存带角度标注图像路径
    功能：
        检测 OBB 并标注置信度最高框旋转角度，保存图片
    """
    img = cv2.imread(image_path)
    if img is None:
        print(f"❌ 无法读取图像：{image_path}")
        return

    model = YOLO(weight_path)
    results = model(img, save=False, imgsz=640, conf=0.15, mode='obb')
    result = results[0]

    boxes = result.obb
    if not boxes:
        print("⚠️ 未检测到目标。")
        return

    best_box = max(boxes, key=lambda x: x.conf.cpu().numpy()[0])
    cx, cy, w, h, r = best_box.xywhr.cpu().numpy()[0]
    angle_deg = np.degrees(r)

    # 绘制 OBB
    annotated_img = img.copy()
    rect = ((cx, cy), (w, h), angle_deg)
    box_pts = cv2.boxPoints(rect).astype(int)
    cv2.polylines(annotated_img, [box_pts], isClosed=True, color=(0, 255, 0), thickness=2)

    # 标注角度
    text = f"{angle_deg:.1f}°"
    font_scale = max(0.5, min(w, h)/100)
    thickness = 2
    text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
    text_x = int(cx - text_size[0]/2)
    text_y = int(cy + text_size[1]/2)
    cv2.putText(annotated_img, text, (text_x, text_y),
                cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255), thickness)

    # 保存
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    cv2.imwrite(save_path, annotated_img)
    print(f"✅ 检测结果已保存至: {save_path}")


# ===============================
# 示例调用
# ===============================
if __name__ == "__main__":
    weight = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb3/weights/best.pt"
    image = r"/home/hx/yolo/output_masks/2.jpg"
    save_path = "./inference_results/best_detected_2.jpg"

    angle_rad = get_best_obb_angle(image, weight)
    print(f"旋转角（弧度）：{angle_rad:.4f}")

    angle_deg = get_best_obb_angle(image, weight, return_degree=True)
    print(f"旋转角（度）：{angle_deg:.2f}°")

    save_obb_visual(image, weight, save_path)