117 lines
4.5 KiB
Python
117 lines
4.5 KiB
Python
import cv2
|
||
import os
|
||
import shutil
|
||
import numpy as np
|
||
from ultralytics import YOLO
|
||
|
||
IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}
|
||
|
||
|
||
def process_obb_images(
|
||
model_path,
|
||
image_dir,
|
||
large_angle_output_dir="./large_angle_images", # 新增：大角度图像保存目录
|
||
conf_thresh=0.15,
|
||
imgsz=640,
|
||
angle_threshold_deg=10.0 # 夹角阈值：超过即视为“大角度”
|
||
):
|
||
"""
|
||
批量处理 OBB 图像，若任意两个目标的主方向夹角 > angle_threshold_deg，
|
||
则将原图移动到 large_angle_output_dir。
|
||
"""
|
||
os.makedirs(large_angle_output_dir, exist_ok=True)
|
||
results_dict = {}
|
||
|
||
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 results_dict
|
||
|
||
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 = []
|
||
has_large_angle = False # 标记是否有大角度的目标
|
||
if boxes is None or len(boxes) == 0:
|
||
print("❌ 该图像中未检测到任何目标")
|
||
else:
|
||
for i, box in enumerate(boxes):
|
||
cls = int(box.cls.cpu().numpy()[0])
|
||
conf = box.conf.cpu().numpy()[0]
|
||
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)
|
||
print(f" Box {i + 1}: Class={cls}, Conf={conf:.3f}, 主方向={angle_deg:.2f}°")
|
||
|
||
# 计算两两夹角
|
||
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_diff_deg = np.degrees(min_diff_rad)
|
||
pairwise_angles_deg.append(angle_diff_deg)
|
||
print(f" Box {i + 1} 与 Box {j + 1} 夹角: {angle_diff_deg:.2f}°")
|
||
if angle_diff_deg > angle_threshold_deg:
|
||
has_large_angle = True
|
||
elif len(angles_deg) == 1:
|
||
print(" 仅检测到一个目标，无法计算夹角")
|
||
|
||
# 保存结果
|
||
results_dict[img_filename] = {
|
||
"angles_deg": angles_deg,
|
||
"pairwise_angles_deg": pairwise_angles_deg,
|
||
"has_large_angle": has_large_angle
|
||
}
|
||
|
||
# 如果存在大角度，移动原图到新文件夹
|
||
if has_large_angle:
|
||
print(f"发现夹角 > {angle_threshold_deg}°，移动原图到大角度文件夹")
|
||
shutil.move(img_path, os.path.join(large_angle_output_dir, img_filename))
|
||
|
||
print(f"\n所有图像处理完成！大角度图像已移动至: {large_angle_output_dir}")
|
||
return results_dict
|
||
|
||
|
||
# ------------------- 测试调用 -------------------
|
||
if __name__ == "__main__":
|
||
MODEL_PATH = r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb_new3/weights/best.pt'
|
||
IMAGE_SOURCE_DIR = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/ready"
|
||
LARGE_ANGLE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/large_angle_images" # 新增输出目录
|
||
|
||
results = process_obb_images(
|
||
model_path=MODEL_PATH,
|
||
image_dir=IMAGE_SOURCE_DIR,
|
||
large_angle_output_dir=LARGE_ANGLE_DIR,
|
||
conf_thresh=0.15,
|
||
imgsz=640,
|
||
angle_threshold_deg=10.0
|
||
)
|
||
|
||
# 可选：打印大角度图像
|
||
large_angle_imgs = [name for name, info in results.items() if info["has_large_angle"]]
|
||
print(f"\n共 {len(large_angle_imgs)} 张图像包含 >10° 的夹角：")
|
||
for name in large_angle_imgs:
|
||
print(f" - {name}") |