85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
from ultralytics import YOLO
|
||
import cv2
|
||
import numpy as np
|
||
import os
|
||
|
||
# 1. 加载模型
|
||
model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb2/weights/best.pt')
|
||
|
||
# 2. 读取图像
|
||
img_path = r"/home/hx/yolo/output_masks/2.jpg"
|
||
img = cv2.imread(img_path)
|
||
|
||
if img is None:
|
||
print(f"❌ 错误：无法读取图像！请检查路径：{img_path}")
|
||
exit(1)
|
||
|
||
# 3. 预测（OBB 模式）
|
||
results = model(
|
||
img,
|
||
save=False,
|
||
imgsz=640,
|
||
conf=0.15,
|
||
mode='obb'
|
||
)
|
||
|
||
# 4. 获取结果并绘制
|
||
result = results[0]
|
||
annotated_img = result.plot()
|
||
|
||
# 5. 保存结果
|
||
output_dir = "./inference_results"
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
filename = os.path.basename(img_path)
|
||
save_path = os.path.join(output_dir, "detected_" + filename)
|
||
cv2.imwrite(save_path, annotated_img)
|
||
print(f"✅ 推理结果已保存至: {save_path}")
|
||
|
||
# 6. 提取旋转框并计算 **两个框之间的夹角**
|
||
boxes = result.obb
|
||
if boxes is None or len(boxes) == 0:
|
||
print("❌ No objects detected.")
|
||
else:
|
||
print(f"✅ Detected {len(boxes)} object(s):")
|
||
directions = [] # 存储每个框的主方向（弧度），归一化到 [0, π)
|
||
|
||
for i, box in enumerate(boxes):
|
||
cls = int(box.cls.cpu().numpy()[0])
|
||
conf = box.conf.cpu().numpy()[0]
|
||
xywhr = box.xywhr.cpu().numpy()[0] # [cx, cy, w, h, r]
|
||
cx, cy, w, h, r_rad = xywhr
|
||
|
||
# 确定主方向（长边方向）
|
||
if w >= h:
|
||
direction = r_rad # 长边方向就是 r
|
||
else:
|
||
direction = r_rad + np.pi / 2 # 长边方向是 r + 90°
|
||
|
||
# 归一化到 [0, π)
|
||
direction = direction % np.pi
|
||
|
||
directions.append(direction)
|
||
angle_deg = np.degrees(direction)
|
||
print(f" Box {i+1}: Class: {cls}, Confidence: {conf:.3f}, 主方向: {angle_deg:.2f}°")
|
||
|
||
# ✅ 计算任意两个框之间的夹角（最小夹角，0° ~ 90°）
|
||
if len(directions) >= 2:
|
||
print("\n🔍 计算两个旋转框之间的夹角（主方向夹角）：")
|
||
for i in range(len(directions)):
|
||
for j in range(i + 1, len(directions)):
|
||
dir1 = directions[i]
|
||
dir2 = directions[j]
|
||
|
||
# 计算方向差（取最小夹角，考虑周期性）
|
||
diff = abs(dir1 - dir2)
|
||
diff = min(diff, np.pi - diff) # 最小夹角（0 ~ π/2）
|
||
diff_deg = np.degrees(diff)
|
||
|
||
print(f" Box {i+1} 与 Box {j+1} 之间的夹角: {diff_deg:.2f}°")
|
||
else:
|
||
print("⚠️ 检测到少于两个目标，无法计算夹角。")
|
||
|
||
# 7. 显示图像
|
||
cv2.imshow("YOLO OBB Prediction", annotated_img)
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows() |