import os
from pathlib import Path
import cv2
from ultralytics import YOLO

# ---------------------------
# 配置路径（请按需修改）
# ---------------------------
MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls/exp_xiantiao_cls/weights/best.pt"          # 你的二分类模型
INPUT_FOLDER = "/home/hx/开发/ML_xiantiao/image/test"        # 输入图像文件夹
OUTPUT_ROOT = "/home/hx/开发/ML_xiantiao/image/test_result"        # 输出根目录（会生成 合格/不合格 子文件夹）

# 类别映射（必须与训练时的 data.yaml 顺序一致）
CLASS_NAMES = {0: "不合格", 1: "合格"}

# ---------------------------
# 批量推理函数
# ---------------------------
def batch_classify(model_path, input_folder, output_root):
    # 加载模型
    model = YOLO(model_path)
    print(f"✅ 模型加载成功: {model_path}")

    # 创建输出目录
    output_root = Path(output_root)
    for cls_name in CLASS_NAMES.values():
        (output_root / cls_name).mkdir(parents=True, exist_ok=True)

    # 支持的图像格式
    IMG_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}

    input_dir = Path(input_folder)
    processed = 0

    for img_path in input_dir.iterdir():
        if img_path.suffix.lower() not in IMG_EXTS:
            continue

        # 读取图像
        img = cv2.imread(str(img_path))
        if img is None:
            print(f"❌ 无法读取: {img_path}")
            continue

        # 推理（整图）
        results = model(img)
        probs = results[0].probs.data.cpu().numpy()
        pred_class_id = int(probs.argmax())
        pred_label = CLASS_NAMES[pred_class_id]
        confidence = float(probs[pred_class_id])

        # 保存原图到对应文件夹
        dst = output_root / pred_label / img_path.name
        cv2.imwrite(str(dst), img)

        print(f"✅ {img_path.name} → {pred_label} ({confidence:.2f})")
        processed += 1

    print(f"\n🎉 共处理 {processed} 张图像，结果已保存至: {output_root}")

# ---------------------------
# 运行入口
# ---------------------------
if __name__ == "__main__":
    batch_classify(
        model_path=MODEL_PATH,
        input_folder=INPUT_FOLDER,
        output_root=OUTPUT_ROOT
    )