from ultralytics import YOLO
import cv2
import os
from pathlib import Path
import time 
# ========== 配置区域 ==========
# 1. 模型路径（替换成你训练好的 best.pt 路径）
MODEL_PATH = 'D:\\deepleaning\\ultralytics-8.4.38\\runs\\detect\\my_train\\exp145\\weights\\best.pt'

# 2. 图片文件夹路径（存放待检测图片的目录）
INPUT_DIR = 'C:\\Users\\123\\Desktop\\text'  # 修改为你的图片文件夹路径

# 3. 检测参数
CONF_THRESHOLD = 0.15   # 置信度阈值，低于此值的检测结果会被过滤
IOU_THRESHOLD = 0.45    # NMS 的 IOU 阈值
IMGSZ = 640             # 推理时的图片尺寸

# 4. 支持的图片格式
IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
# =============================

def batch_predict():
    # 加载模型
    print(f"正在加载模型: {MODEL_PATH}")
    model = YOLO(MODEL_PATH)
    
    # 创建结果输出目录
    input_path = Path(INPUT_DIR)
    output_dir = input_path / 'result'
    output_dir.mkdir(exist_ok=True)
    print(f"结果将保存到: {output_dir}")
    
    # 获取所有图片文件
    image_files = [f for f in input_path.iterdir() 
                   if f.suffix.lower() in IMAGE_EXTENSIONS]
    
    if not image_files:
        print(f"未在 {INPUT_DIR} 中找到支持的图片文件")
        print(f"支持的格式: {', '.join(IMAGE_EXTENSIONS)}")
        return
    
    print(f"找到 {len(image_files)} 张图片，开始检测...")
    
    # 统计信息
    total_objects = 0
    images_with_objects = 0
    inference_time = 0.0  # 单张图片推理时间（秒）
    total_inference_time = 0.0  # 累计推理时间（秒）
    
    # 逐张处理
    for idx, img_path in enumerate(image_files, 1):
        print(f"[{idx}/{len(image_files)}] 正在处理: {img_path.name}")
        
        # 🕒 开始计时
        start_time = time.time()

        # 推理
        results = model(img_path, conf=CONF_THRESHOLD, iou=IOU_THRESHOLD, imgsz=IMGSZ)
        
        # 🕒 结束计时
        end_time = time.time()
        inference_time = end_time - start_time
        total_inference_time += inference_time
        

        # 获取检测结果
        boxes = results[0].boxes
        num_objects = len(boxes) if boxes is not None else 0
        total_objects += num_objects
        if num_objects > 0:
            images_with_objects += 1
        
        # 保存标注后的图片
        output_path = output_dir / img_path.name
        # 绘制检测框并保存（直接用 results 的 plot 方法）
        annotated_img = results[0].plot()  # 返回 BGR 格式的 numpy 数组
        cv2.imwrite(str(output_path), annotated_img)

        # 打印该图片的检测详情
        print(f"    ⏱️ 检测耗时: {inference_time*1000:.2f} ms")
        if boxes is not None:
            for box in boxes:
                conf = float(box.conf[0])
                cls_id = int(box.cls[0])
                cls_name = model.names[cls_id]
                print(f"    → 检出: {cls_name}, 置信度: {conf:.4f}")
    
    # 输出汇总统计（增加时间信息）
    print("\n" + "="*50)
    print("检测完成！汇总统计：")
    print(f"  总图片数: {len(image_files)}")
    print(f"  检出目标的图片数: {images_with_objects}")
    print(f"  总检出目标数: {total_objects}")
    print(f"  总推理时间: {total_inference_time:.2f} 秒")
    print(f"  平均每张图片耗时: {total_inference_time/len(image_files)*1000:.2f} ms")
    print(f"  结果保存位置: {output_dir}")
    print("="*50)

if __name__ == '__main__':
    batch_predict()