from ultralytics import YOLO
import cv2

# 1. 加载模型
model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb/weights/best.pt')

# 2. 读取图像
img_path = r"/home/hx/yolo/test_image/2.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,      # 必须和训练时一致（你训练日志中是 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}")