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

# ------------------ 模型与路径配置 ------------------
MODEL_PATH = '../ultralytics_yolo11-main/runs/train/seg_j/exp/weights/best.pt'
OUTPUT_DIR = '../test_image'
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 加载模型（全局一次）
model = YOLO(MODEL_PATH)
model.to('cuda')  # 如无 GPU，可改为 'cpu'


def detect_jaw_angle(image_path, mode='show'):
    """
    检测图像中两个夹具的开合角度

    参数:
        image_path (str): 输入图像路径
        mode (str): 'show' -> 保存可视化结果；'silent' -> 只返回角度

    返回:
        float: 夹具开合角度 (0 ~ 90°)
    """
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {image_path}")

    h, w = img.shape[:2]
    filename = os.path.basename(image_path)
    name_only = os.path.splitext(filename)[0]

    # 创建掩码并检测
    composite_mask = np.zeros((h, w), dtype=np.uint8)
    results = model(image_path, imgsz=640, conf=0.5)

    jaws = []
    for r in results:
        if r.masks is not None:
            masks = r.masks.data.cpu().numpy()
            boxes = r.boxes.xyxy.cpu().numpy()

            for i, mask in enumerate(masks):
                x1, y1, x2, y2 = map(int, boxes[i])
                x1, y1 = max(0, x1), max(0, y1)
                x2, y2 = min(w, x2), min(h, y2)

                obj_mask = np.zeros((h, w), dtype=np.uint8)
                mask_resized = cv2.resize(mask, (w, h))
                obj_mask[y1:y2, x1:x2] = (mask_resized[y1:y2, x1:x2] * 255).astype(np.uint8)

                contours, _ = cv2.findContours(obj_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
                if len(contours) == 0:
                    continue

                largest_contour = max(contours, key=cv2.contourArea)
                area = cv2.contourArea(largest_contour)
                if area < 100:
                    continue

                rect = cv2.minAreaRect(largest_contour)
                jaws.append({
                    'contour': largest_contour,
                    'rect': rect,
                    'area': area
                })
                composite_mask = np.maximum(composite_mask, obj_mask)

    # 必须检测到至少两个夹具
    if len(jaws) < 2:
        print(f"检测到的夹具少于2个（共{len(jaws)}个）")
        if mode == 'show':
            cv2.imwrite(os.path.join(OUTPUT_DIR, f'mask_{name_only}.png'), composite_mask)
        return None

    # 取面积最大的两个
    jaws.sort(key=lambda x: x['area'], reverse=True)
    jaw1, jaw2 = jaws[0], jaws[1]

    # ------------------ 计算角度 ------------------
    def get_long_edge_vector(rect):
        center, (w, h), angle = rect
        rad = np.radians(angle + (0 if w >= h else 90))
        return np.array([np.cos(rad), np.sin(rad)])

    def get_center(contour):
        M = cv2.moments(contour)
        return np.array([M['m10']/M['m00'], M['m01']/M['m00']]) if M['m00'] != 0 else np.array([0,0])

    dir1 = get_long_edge_vector(jaw1['rect'])
    dir2 = get_long_edge_vector(jaw2['rect'])

    # 校正方向：指向夹具中心
    center1 = get_center(jaw1['contour'])
    center2 = get_center(jaw2['contour'])
    fixture_center = (center1 + center2) / 2

    to_center1 = fixture_center - center1
    to_center2 = fixture_center - center2

    if np.linalg.norm(to_center1) > 1e-6 and np.dot(dir1, to_center1) < 0:
        dir1 = -dir1
    if np.linalg.norm(to_center2) > 1e-6 and np.dot(dir2, to_center2) < 0:
        dir2 = -dir2

    # 计算最小内角
    cos_angle = np.clip(np.dot(dir1, dir2), -1.0, 1.0)
    angle = np.degrees(np.arccos(cos_angle))
    opening_angle = min(angle, 180 - angle)

    # ------------------ 可视化（仅 mode='show'）------------------
    if mode == 'show':
        vis_img = np.stack([composite_mask]*3, axis=-1)
        vis_img[composite_mask > 0] = [255, 255, 255]

        # 绘制矩形框
        box1 = np.int32(cv2.boxPoints(jaw1['rect']))
        box2 = np.int32(cv2.boxPoints(jaw2['rect']))
        cv2.drawContours(vis_img, [box1], 0, (0, 0, 255), 2)
        cv2.drawContours(vis_img, [box2], 0, (255, 0, 0), 2)

        # 绘制方向箭头
        scale = 60
        c1 = tuple(np.int32(center1))
        c2 = tuple(np.int32(center2))
        end1 = tuple(np.int32(center1 + scale * dir1))
        end2 = tuple(np.int32(center2 + scale * dir2))
        cv2.arrowedLine(vis_img, c1, end1, (0, 255, 0), 2, tipLength=0.3)
        cv2.arrowedLine(vis_img, c2, end2, (0, 255, 0), 2, tipLength=0.3)

        # 标注角度
        cv2.putText(vis_img, f"Angle: {opening_angle:.2f}°", (20, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

        # 保存
        save_path = os.path.join(OUTPUT_DIR, f'mask_with_angle_{name_only}.png')
        cv2.imwrite(save_path, vis_img)
        print(f"✅ 可视化结果已保存: {save_path}")

    return round(opening_angle, 2)

# ------------------ 主函数 ------------------
if __name__ == '__main__':
    # ✅ 设置输入图像路
    image_path = r"/home/hx/yolo/output_masks/2.jpg"  # ← 修改为你自己的图片路径

    # ✅ 模式选择：
    #   mode='show': 保存可视化图像
    #   mode='silent': 只返回角度
    mode = 'show'  # 或 'silent'

    print(f"🔍 正在处理图像: {image_path}")
    angle = detect_jaw_angle(image_path, mode=mode)

    if angle is not None:
        print(f"🎉 检测到的夹具开合角度: {angle}°")
    else:
        print("❌ 未能检测到足够的夹具")