153 lines
4.6 KiB
Python
153 lines
4.6 KiB
Python
import cv2
|
|
import os
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
|
|
IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}
|
|
|
|
|
|
def draw_direction(img, cx, cy, angle_deg, length=80, color=(0, 255, 0)):
|
|
"""画主方向箭头"""
|
|
theta = np.radians(angle_deg)
|
|
x2 = int(cx + length * np.cos(theta))
|
|
y2 = int(cy + length * np.sin(theta))
|
|
cv2.arrowedLine(
|
|
img,
|
|
(int(cx), int(cy)),
|
|
(x2, y2),
|
|
color,
|
|
2,
|
|
tipLength=0.2
|
|
)
|
|
|
|
|
|
def process_obb_images(
|
|
model_path,
|
|
image_dir,
|
|
output_dir="./inference_results",
|
|
conf_thresh=0.15,
|
|
imgsz=640
|
|
):
|
|
os.makedirs(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("❌ 读取失败，跳过")
|
|
continue
|
|
|
|
# ---------- OBB 推理 ----------
|
|
results = model(img, save=False, imgsz=imgsz, conf=conf_thresh, mode="obb")
|
|
result = results[0]
|
|
annotated_img = result.plot()
|
|
|
|
boxes = result.obb
|
|
angles_deg = []
|
|
centers = []
|
|
|
|
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)
|
|
centers.append((int(cx), int(cy)))
|
|
|
|
print(f" Box {i + 1} 主方向: {angle_deg:.2f} deg")
|
|
|
|
# 主方向可视化
|
|
draw_direction(annotated_img, cx, cy, angle_deg)
|
|
cv2.circle(annotated_img, (int(cx), int(cy)), 4, (0, 0, 255), -1)
|
|
|
|
# ---------- 两两夹角 ----------
|
|
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_ij = np.degrees(min_diff_rad)
|
|
|
|
pairwise_angles_deg.append(angle_ij)
|
|
print(
|
|
f" Box {i + 1} 与 Box {j + 1} 夹角: {angle_ij:.2f} deg"
|
|
)
|
|
|
|
# ---------- 右上角粗体 angle ----------
|
|
if pairwise_angles_deg:
|
|
max_angle = max(pairwise_angles_deg)
|
|
h, w = annotated_img.shape[:2]
|
|
|
|
text = f"angle: {max_angle:.1f} deg"
|
|
|
|
# 粗体效果（多次叠加）
|
|
for dx, dy in [(0, 0), (1, 0), (0, 1), (1, 1)]:
|
|
cv2.putText(
|
|
annotated_img,
|
|
text,
|
|
(w - 300 + dx, 40 + dy),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
1.2,
|
|
(0, 0, 255),
|
|
3
|
|
)
|
|
|
|
# ---------- 保存 ----------
|
|
save_path = os.path.join(output_dir, "detected_" + img_filename)
|
|
cv2.imwrite(save_path, annotated_img)
|
|
print(f"✅ 保存完成: {save_path}")
|
|
|
|
results_dict[img_filename] = {
|
|
"angles_deg": angles_deg,
|
|
"pairwise_angles_deg": pairwise_angles_deg
|
|
}
|
|
|
|
print("\n🎉 全部处理完成")
|
|
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"/home/hx/yolo/angle_base_obb/test_image"
|
|
OUTPUT_DIR = "./inference_results"
|
|
|
|
results = process_obb_images(
|
|
MODEL_PATH,
|
|
IMAGE_SOURCE_DIR,
|
|
OUTPUT_DIR
|
|
)
|
|
|
|
for img_name, info in results.items():
|
|
print(f"\n{img_name}")
|
|
print("主方向角:", info["angles_deg"])
|
|
print("两两夹角:", info["pairwise_angles_deg"])
|