zjsh_yolov11/angle_base_seg/test_seg_angle_f1.py

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

# ------------------ 配置 ------------------
model_path = '/home/hx/yolo/ultralytics_yolo11-main/runs/train/seg_j/exp/weights/best.pt'
img_folder = '/home/hx/yolo/output_masks'
output_mask_dir = 'output_masks1'
os.makedirs(output_mask_dir, exist_ok=True)

SUPPORTED_FORMATS = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff')

# ------------------ 加载模型 ------------------
model = YOLO(model_path)
model.to('cuda')  # 使用 GPU 加速（如有））

def get_orientation_vector(rect):
    """从 minAreaRect 中提取主方向向量（单位向量）"""
    center, (width, height), angle = rect
    if width >= height:
        rad = np.radians(angle)
    else:
        rad = np.radians(angle + 90)
    direction = np.array([np.cos(rad), np.sin(rad)])
    return direction / (np.linalg.norm(direction) + 1e-8)

def get_contour_center(contour):
    """计算轮廓质心"""
    M = cv2.moments(contour)
    if M["m00"] == 0:
        return np.array([0, 0])
    return np.array([int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"])])

def calculate_jaw_opening_angle(jaw1, jaw2):
    """
    计算夹具开合角度，并返回修正后的方向向量用于可视化
    返回: (angle, dir1_final, dir2_final)
    """
    center1 = get_contour_center(jaw1['contour'])
    center2 = get_contour_center(jaw2['contour'])
    fixture_center = np.array([(center1[0] + center2[0]) / 2.0, (center1[1] + center2[1]) / 2.0])

    # 获取原始主方向
    dir1_orig = get_orientation_vector(jaw1['rect'])
    dir2_orig = get_orientation_vector(jaw2['rect'])

    def correct_and_compute(d1, d2):
        """校正方向并计算夹角"""
        # 校正：确保方向指向夹具中心
        to_center1 = fixture_center - center1
        if np.linalg.norm(to_center1) > 1e-6:
            to_center1 = to_center1 / np.linalg.norm(to_center1)
            if np.dot(d1, to_center1) < 0:
                d1 = -d1

        to_center2 = fixture_center - center2
        if np.linalg.norm(to_center2) > 1e-6:
            to_center2 = to_center2 / np.linalg.norm(to_center2)
            if np.dot(d2, to_center2) < 0:
                d2 = -d2

        # 计算夹角
        cos_angle = np.clip(np.dot(d1, d2), -1.0, 1.0)
        angle = np.degrees(np.arccos(cos_angle))
        return angle, d1, d2

    # --- 第一步：原始方向计算 ---
    angle_raw, dir1_raw, dir2_raw = correct_and_compute(dir1_orig, dir2_orig)

    if angle_raw <= 170.0:
        print(f"✅ 角度正常: {angle_raw:.2f}°")
        return angle_raw, dir1_raw, dir2_raw

    print(f"⚠️  初始角度过大: {angle_raw:.2f}°，尝试翻转 jaw2 主方向...")

    # --- 第二步：jaw2 主方向取反后重新计算 ---
    angle_corrected, dir1_corr, dir2_corr = correct_and_compute(dir1_orig, -dir2_orig)
    print(f"🔄 方向修正后: {angle_corrected:.2f}°")

    # --- 第三步：数值兜底修正（角度仍过大）---
    if angle_corrected > 170.0:
        final_angle = 180.0 - angle_corrected
        print(f"🔧 数值修正: {angle_corrected:.2f}° → {final_angle:.2f}°")
        return final_angle, dir1_corr, dir2_corr

    return angle_corrected, dir1_corr, dir2_corr

def process_image(img_path, output_dir):
    img = cv2.imread(img_path)
    if img is None:
        print(f"❌ 无法读取图像: {img_path}")
        return

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

    print(f"\n🔄 正在处理: {filename}")

    # 初始化合并掩码
    composite_mask = np.zeros((h, w), dtype=np.uint8)
    results = model(img_path, imgsz=1280, conf=0.5)

    rotated_rects = []

    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)
                rotated_rects.append({
                    'rect': rect,
                    'contour': largest_contour,
                    'area': area
                })

                composite_mask = np.maximum(composite_mask, obj_mask)

    # 保存合并掩码
    #mask_save_path = os.path.join(output_dir, f'mask_{name_only}.png')
    #cv2.imwrite(mask_save_path, composite_mask)
    #print(f"✅ 掩码已保存: {mask_save_path}")

    if len(rotated_rects) < 2:
        print(f"⚠️ 检测到的对象少于2个（共{len(rotated_rects)}个）: {filename}")
        return

    # 按面积排序，取前两个
    rotated_rects.sort(key=lambda x: x['area'], reverse=True)
    jaw1, jaw2 = rotated_rects[0], rotated_rects[1]

    # ------------------ 计算角度 + 获取修正后的方向 ------------------
    opening_angle, dir1_final, dir2_final = calculate_jaw_opening_angle(jaw1, jaw2)
    print(f"✅ 最终夹具开合角度: {opening_angle:.2f}°")

    # ------------------ 可视化 ------------------
    vis_img = img.copy()

    # 绘制最小外接矩形
    box1 = cv2.boxPoints(jaw1['rect'])
    box1 = np.int32(box1)
    cv2.drawContours(vis_img, [box1], 0, (0, 0, 255), 2)  # jaw1: 红色

    box2 = cv2.boxPoints(jaw2['rect'])
    box2 = np.int32(box2)
    cv2.drawContours(vis_img, [box2], 0, (255, 0, 0), 2)  # jaw2: 蓝色

    # 计算中心点
    center1 = get_contour_center(jaw1['contour'])
    center2 = get_contour_center(jaw2['contour'])
    fixture_center = ((center1[0] + center2[0]) // 2, (center1[1] + center2[1]) // 2)

    # ✅ 使用修正后的方向向量绘制箭头
    scale = 50
    end1 = center1 + scale * dir1_final
    end2 = center2 + scale * dir2_final

    cv2.arrowedLine(vis_img, tuple(center1), tuple(end1.astype(int)), (0, 255, 0), 2, tipLength=0.3)  # 绿色箭头
    cv2.arrowedLine(vis_img, tuple(center2), tuple(end2.astype(int)), (0, 255, 0), 2, tipLength=0.3)

    # 标注夹具中心
    cv2.circle(vis_img, fixture_center, 5, (255, 255, 0), -1)
    cv2.putText(vis_img, "Fixture Center", (fixture_center[0] + 10, fixture_center[1]),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)

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

    # 保存可视化图像
    vis_save_path = os.path.join(output_dir, f'angle_{name_only}.jpg')
    cv2.imwrite(vis_save_path, vis_img)
    print(f"✅ 可视化图像已保存: {vis_save_path}")

# ------------------ 主程序 ------------------
if __name__ == '__main__':
    if not os.path.isdir(img_folder):
        print(f"❌ 图像文件夹不存在: {img_folder}")
        exit()

    image_files = [f for f in os.listdir(img_folder) if f.lower().endswith(SUPPORTED_FORMATS)]
    if len(image_files) == 0:
        print(f"⚠️ 未找到支持的图像文件")
        exit()

    print(f"✅ 发现 {len(image_files)} 张图像，开始处理...")
    for image_file in image_files:
        image_path = os.path.join(img_folder, image_file)
        process_image(image_path, output_mask_dir)

    print(f"\n🎉 所有图像处理完成！结果保存在: {output_mask_dir}")