
import os
from ultralytics import YOLO
import argparse
from model_utils import get_project_root, find_latest_model


def predict_image(model_path=None, source=None, conf=0.25, save=True):
    """预测单张图片或目录
    
    Args:
        model_path: 模型路径（如果为 None 则自动查找）
        source: 预测源（图片/目录）
        conf: 置信度阈值
        save: 是否保存结果
    """
    project_root = get_project_root()
    
    # 如果没提供模型路径，自动查找最新模型
    if model_path is None or not os.path.exists(model_path):
        found_model = find_latest_model()
        if found_model:
            model_path = found_model
            print(f"找到最新模型: {model_path}")
        else:
            print(f"\n错误: 找不到模型文件！")
            print(f"\n可能的原因:")
            print(f"  1. 还没有运行训练，先运行: python src/train.py")
            print(f"  2. 使用 --model 参数指定正确路径")
            return None
    else:
        print(f"模型路径: {model_path}")
    
    # 如果没提供 source，使用配置文件中的 test 目录
    if source is None:
        config_path = os.path.join(project_root, 'dataset.yaml')
        if os.path.exists(config_path):
            import yaml
            with open(config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            if 'path' in config and 'test' in config:
                source = os.path.join(config['path'], config['test'])
    
    if source is None or not os.path.exists(source):
        print(f"\n警告: 找不到预测源，请使用 --source 参数指定图片路径")
    
    try:
        # 加载模型
        model = YOLO(model_path)
        
        # 预测
        results = model.predict(source, save=save, conf=conf)
        
        # 打印结果
        for result in results:
            boxes = result.boxes
            print(f"\n图片: {result.path}")
            if len(boxes) == 0:
                print(f"  未检测到目标")
            else:
                for i in range(len(boxes)):
                    class_name = model.names[int(boxes.cls[i])]
                    confidence = float(boxes.conf[i])
                    print(f"  检测到: {class_name}, 置信度: {confidence:.2f}")
        
        return results
        
    except Exception as e:
        print(f"\n预测失败: {e}")
        return None


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='YOLO 预测')
    parser.add_argument('--model', default=None,
                        help='模型路径（默认：自动查找最新模型）')
    parser.add_argument('--source', default=None,
                        help='预测源（图片/目录，默认：读取 dataset.yaml）')
    parser.add_argument('--conf', type=float, default=0.25,
                        help='置信度阈值（默认: 0.25）')
    parser.add_argument('--nosave', action='store_true',
                        help='不保存结果（默认保存）')
    
    args = parser.parse_args()
    
    predict_image(args.model, args.source, conf=args.conf, save=not args.nosave)
