74 lines
2.5 KiB
Python
74 lines
2.5 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_obb/weights/best.pt')
|
||
|
||
# 2. 读取图像
|
||
img_path = r"/output_masks/3.png"
|
||
img = cv2.imread(img_path)
|
||
|
||
# ✅ 检查图像是否加载成功
|
||
if img is None:
|
||
print(f"❌ 错误：无法读取图像！请检查路径：{img_path}")
|
||
print("💡 提示：可能是文件不存在、路径错误或图像损坏")
|
||
exit(1)
|
||
|
||
# 3. 预测（使用 OBB 模式！）
|
||
results = model(
|
||
img,
|
||
save=False,
|
||
imgsz=1280, # 必须和训练时一致
|
||
conf=0.25,
|
||
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):")
|
||
angles = [] # 存储每个框的旋转角度（角度制）
|
||
|
||
for i, box in enumerate(boxes):
|
||
cls = int(box.cls.cpu().numpy()[0])
|
||
conf = box.conf.cpu().numpy()[0]
|
||
|
||
# ✅ 正确方式：使用 .xywhr 获取旋转角度（新版本 API）
|
||
xywhr = box.xywhr.cpu().numpy()[0] # [x_center, y_center, width, height, rotation]
|
||
angle_rad = xywhr[4] # 第5个值是旋转角度（弧度）
|
||
angle_deg = np.degrees(angle_rad) # 转为角度
|
||
|
||
angles.append(angle_deg)
|
||
print(f" Box {i+1}: Class: {cls}, Confidence: {conf:.3f}, Angle: {angle_deg:.2f}°")
|
||
|
||
# ✅ 计算任意两个框之间的最小夹角差
|
||
if len(angles) >= 2:
|
||
print("\n🔍 计算旋转框之间的夹角差：")
|
||
for i in range(len(angles)):
|
||
for j in range(i + 1, len(angles)):
|
||
diff = abs(angles[i] - angles[j])
|
||
min_angle_diff = min(diff, 180 - diff) # 取最小夹角（0~180°内）
|
||
print(f" Box {i+1} 与 Box {j+1} 的最小夹角差: {min_angle_diff:.2f}°")
|
||
else:
|
||
print("⚠️ 检测到少于两个目标，无法计算夹角。")
|
||
|
||
# ✅ 7. 显示图像
|
||
cv2.imshow("YOLO OBB Prediction", annotated_img)
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows() |