2025-08-18 15:13:54 +08:00
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
|
import cv2
|
|
|
|
|
|
|
|
|
|
|
|
# 1. 加载模型
|
2025-09-01 14:14:18 +08:00
|
|
|
|
model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb/weights/best.pt')
|
2025-08-18 15:13:54 +08:00
|
|
|
|
|
|
|
|
|
|
# 2. 读取图像
|
2025-09-01 14:14:18 +08:00
|
|
|
|
img_path = r"/home/hx/yolo/test_image/2.png"
|
2025-08-18 15:13:54 +08:00
|
|
|
|
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, # 必须和训练时一致（你训练日志中是 1280）
|
|
|
|
|
|
conf=0.25,
|
|
|
|
|
|
mode='obb' # 👈👈👈 关键！必须加这一行才能启用旋转框模式
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
# 4. 在原图上绘制检测结果
|
|
|
|
|
|
result = results[0]
|
|
|
|
|
|
annotated_img = result.plot() # OBB 模式下会自动画旋转框
|
|
|
|
|
|
|
|
|
|
|
|
# 5. 显示图像
|
|
|
|
|
|
cv2.imshow("YOLO OBB Prediction", annotated_img)
|
|
|
|
|
|
cv2.waitKey(0)
|
|
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
|
|
|
|
|
|
# ✅ 可选：打印检测结果
|
|
|
|
|
|
if result.boxes is None or len(result.boxes) == 0:
|
|
|
|
|
|
print("❌ No objects detected.")
|
|
|
|
|
|
else:
|
|
|
|
|
|
print(f"✅ Detected {len(result.boxes)} object(s):")
|
|
|
|
|
|
for box in result.boxes:
|
|
|
|
|
|
cls = int(box.cls.cpu().numpy()[0])
|
|
|
|
|
|
conf = box.conf.cpu().numpy()[0]
|
|
|
|
|
|
print(f" Class: {cls}, Confidence: {conf:.3f}")
|
