43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
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}") |