from ultralytics import YOLO
import os

# 确保 if __name__ == '__main__': 结构，这是一个好习惯
if __name__ == '__main__':
    # 1. 加载你训练好的最佳模型
    # !! 修改为你自己的 best.pt 路径 !!
    model_path = r'runs_up/detect/train/weights/best.pt'
    model = YOLO(model_path)

    # 2. 指定你要预测的图片或文件夹
    # 可以是单张图片路径，也可以是整个文件夹的路径
    # 强烈建议使用验证集里的图片，或者一些全新的测试图片
    source_path = r'train_data_up/images/val' # 预测整个验证集文件夹

    # 3. 执行预测
    # save=True:  会将画好框的图片保存下来
    # conf=0.5:   只显示置信度大于 0.5 的预测结果，可以调整这个值
    results = model.predict(source=source_path, save=True, conf=0.5)

    # 预测结果会默认保存在 runs_up/detect/predictX 文件夹下
    print("\n预测完成！")
    # 你可以从 results 对象中获取详细信息，但对于可视化验证，直接去看保存的图片更方便。