106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
|
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 推理，计算目标主方向夹角分布（每10°一个区间），
|
|||
|
|
并统计最大和最小夹角。
|
|||
|
|
|
|||
|
|
输入：
|
|||
|
|
model_path: YOLO 权重路径
|
|||
|
|
image_dir: 图像文件夹路径
|
|||
|
|
conf_thresh: 置信度阈值
|
|||
|
|
imgsz: 输入图像大小
|
|||
|
|
输出：
|
|||
|
|
distribution: 每10°一个区间的夹角统计字典
|
|||
|
|
max_angle: 所有图像的最大夹角
|
|||
|
|
min_angle: 所有图像的最小夹角
|
|||
|
|
"""
|
|||
|
|
# 初始化每10°统计字典
|
|||
|
|
angle_distribution = {f"{i}-{i + 10}": 0 for i in range(0, 180, 10)}
|
|||
|
|
max_angle = -1
|
|||
|
|
min_angle = 181
|
|||
|
|
|
|||
|
|
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, max_angle, min_angle
|
|||
|
|
|
|||
|
|
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("❌ 该图像中未检测到任何目标")
|
|||
|
|
continue
|
|||
|
|
else:
|
|||
|
|
for box in 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)
|
|||
|
|
|
|||
|
|
# 两两夹角
|
|||
|
|
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_diff_deg = np.degrees(min_diff_rad)
|
|||
|
|
|
|||
|
|
# 更新每10°分布
|
|||
|
|
bin_index = int(angle_diff_deg // 10) * 10
|
|||
|
|
bin_key = f"{bin_index}-{bin_index + 10}"
|
|||
|
|
if bin_key in angle_distribution:
|
|||
|
|
angle_distribution[bin_key] += 1
|
|||
|
|
|
|||
|
|
# 更新全局最大最小夹角
|
|||
|
|
if angle_diff_deg > max_angle:
|
|||
|
|
max_angle = angle_diff_deg
|
|||
|
|
if angle_diff_deg < min_angle:
|
|||
|
|
min_angle = angle_diff_deg
|
|||
|
|
|
|||
|
|
print("\n所有图像处理完成！")
|
|||
|
|
return angle_distribution, max_angle, min_angle
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ------------------- 测试调用 -------------------
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
MODEL_PATH = r'obb.pt'
|
|||
|
|
IMAGE_SOURCE_DIR = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb3/train"
|
|||
|
|
|
|||
|
|
distribution, max_angle, min_angle = process_obb_images_for_angle_distribution(MODEL_PATH, IMAGE_SOURCE_DIR)
|
|||
|
|
|
|||
|
|
print("\n夹角分布统计（每10°）:")
|
|||
|
|
for k, v in distribution.items():
|
|||
|
|
print(f"{k}°: {v}")
|
|||
|
|
|
|||
|
|
print(f"\n最大夹角: {max_angle:.2f}°")
|
|||
|
|
print(f"最小夹角: {min_angle:.2f}°")
|
|||
|
|
|