zjsh_yolov11/angle_base_obb/tongji.py

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

IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}


def process_obb_images_for_angle_distribution(model_path, image_dir, conf_thresh=0.15, imgsz=640):
    """
    批量处理图像的 OBB 推理，计算每张图像检测目标的主方向和夹角，并统计夹角分布情况。

    输入：
        model_path: YOLO 权重路径
        image_dir: 图像文件夹路径
        conf_thresh: 置信度阈值
        imgsz: 输入图像大小
    输出：
        angle_distribution: {'<6': count, '6-20': count, '>20': count}
    """
    results_dict = {}
    angle_distribution = {'<6': 0, '6-20': 0, '>20': 0}

    print("加载 YOLO 模型...")
    model = YOLO(model_path)
    print("✅ 模型加载完成")

    # 获取图像文件
    image_files = [f for f in os.listdir(image_dir) if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS]
    if not image_files:
        print(f"❌ 未找到图像文件：{image_dir}")
        return angle_distribution

    print(f"发现 {len(image_files)} 张图像待处理")

    for img_filename in image_files:
        img_path = os.path.join(image_dir, img_filename)
        print(f"\n正在处理：{img_filename}")

        img = cv2.imread(img_path)
        if img is None:
            print(f"❌ 跳过：无法读取图像 {img_path}")
            continue

        # 推理 OBB
        results = model(img, save=False, imgsz=imgsz, conf=conf_thresh, mode='obb')
        result = results[0]

        # 提取旋转角
        boxes = result.obb
        angles_deg = []
        if boxes is None or len(boxes) == 0:
            print("❌ 该图像中未检测到任何目标")
        else:
            for i, box in enumerate(boxes):
                cx, cy, w, h, r_rad = box.xywhr.cpu().numpy()[0]
                direction = r_rad if w >= h else r_rad + np.pi / 2
                direction = direction % np.pi
                angle_deg = np.degrees(direction)
                angles_deg.append(angle_deg)

        # 两两夹角
        pairwise_angles_deg = []
        if len(angles_deg) >= 2:
            for i in range(len(angles_deg)):
                for j in range(i + 1, len(angles_deg)):
                    diff_rad = abs(np.radians(angles_deg[i]) - np.radians(angles_deg[j]))
                    min_diff_rad = min(diff_rad, np.pi - diff_rad)
                    angle_deg_diff = np.degrees(min_diff_rad)
                    pairwise_angles_deg.append(angle_deg_diff)

                    # 更新角度分布统计
                    if angle_deg_diff < 6:
                        angle_distribution['<6'] += 1
                    elif 6 <= angle_deg_diff <= 20:
                        angle_distribution['6-20'] += 1
                    else:
                        angle_distribution['>20'] += 1

                    print(f"  Box {i + 1} 与 Box {j + 1} 夹角: {angle_deg_diff:.2f}°")

        # 保存每张图像结果
        results_dict[img_filename] = {
            "angles_deg": angles_deg,
            "pairwise_angles_deg": pairwise_angles_deg
        }

    print("\n所有图像处理完成！")
    return angle_distribution


# ------------------- 测试调用 -------------------
if __name__ == "__main__":
    MODEL_PATH = r'obb.pt'
    IMAGE_SOURCE_DIR = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb3/train"

    distribution = process_obb_images_for_angle_distribution(MODEL_PATH, IMAGE_SOURCE_DIR)
    print("\n夹角分布统计:")
    print(f"小于6度的夹角数量: {distribution['<6']}")
    print(f"在6至20度之间的夹角数量: {distribution['6-20']}")
    print(f"大于20度的夹角数量: {distribution['>20']}")