Files
zjsh_yolov11/angle_base_obb/test.py

43 lines
1.3 KiB
Python
Raw Normal View History

2025-08-18 15:13:54 +08:00
from ultralytics import YOLO
import cv2
# 1. 加载模型
model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb2/weights/best.pt')
# 2. 读取图像
img_path = r"/home/hx/桌面/image/images/test/1.jpg"
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}")