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

IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}


class ObjectDetector:
    """封装 YOLO 目标检测模型"""

    def __init__(self, model_path):
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        self.model = YOLO(model_path)
        print(f"[INFO] 成功加载 YOLO 目标检测模型: {model_path}")

    def detect(self, img_np, conf_threshold=0.0):
        """返回所有置信度 >= conf_threshold 的检测结果"""
        results = self.model.predict(img_np, conf=conf_threshold, verbose=False)
        detections = []
        for result in results:
            boxes = result.boxes.cpu().numpy()
            for box in boxes:
                detection_info = {
                    'bbox_xyxy': box.xyxy[0],  # [x1, y1, x2, y2]
                    'confidence': float(box.conf.item()),
                    'class_id': int(box.cls.item())
                }
                detections.append(detection_info)
        return detections


def save_yolo_detect_labels_from_folder(
        model_path,
        image_dir,
        output_dir,
        conf_threshold=0.5,
        label_map={0: "hole", 1: "crack"}  # 可选，仅用于日志
):
    """
    对 image_dir 中所有图像进行 YOLO Detect 推理，
    每个类别保留最高置信度框，保存为 YOLO 格式的 .txt 标签文件。

    YOLO 格式: <class_id> <cx_norm> <cy_norm> <w_norm> <h_norm>
    """
    image_dir = Path(image_dir)
    output_dir = Path(output_dir)
    labels_dir = output_dir / "labels"
    labels_dir.mkdir(parents=True, exist_ok=True)

    # 获取图像列表
    image_files = [
        f for f in sorted(os.listdir(image_dir))
        if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS
    ]
    if not image_files:
        print(f"❌ 未在 {image_dir} 中找到支持的图像文件")
        return

    print(f"共找到 {len(image_files)} 张图像，开始推理...")
    detector = ObjectDetector(model_path)

    for img_filename in image_files:
        img_path = image_dir / img_filename
        stem = Path(img_filename).stem
        txt_path = labels_dir / f"{stem}.txt"

        # 读图
        img = cv2.imread(str(img_path))
        if img is None:
            print(f"⚠️ 跳过无效图像: {img_path}")
            txt_path.write_text("")  # 写空文件
            continue

        H, W = img.shape[:2]

        # 推理（获取所有 ≥ conf_threshold 的框）
        all_detections = detector.detect(img, conf_threshold=conf_threshold)

        # 按类别保留最高置信度框
        best_per_class = {}
        for det in all_detections:
            cls_id = det['class_id']
            if cls_id not in best_per_class or det['confidence'] > best_per_class[cls_id]['confidence']:
                best_per_class[cls_id] = det

        top_detections = list(best_per_class.values())

        # 转为 YOLO 格式并写入
        lines = []
        for det in top_detections:
            x1, y1, x2, y2 = det['bbox_xyxy']
            cx = (x1 + x2) / 2.0
            cy = (y1 + y2) / 2.0
            bw = x2 - x1
            bh = y2 - y1

            # 归一化
            cx_norm = cx / W
            cy_norm = cy / H
            w_norm = bw / W
            h_norm = bh / H

            # 限制在 [0, 1]
            cx_norm = max(0.0, min(1.0, cx_norm))
            cy_norm = max(0.0, min(1.0, cy_norm))
            w_norm = max(0.0, min(1.0, w_norm))
            h_norm = max(0.0, min(1.0, h_norm))

            line = f"{det['class_id']} {cx_norm:.6f} {cy_norm:.6f} {w_norm:.6f} {h_norm:.6f}"
            lines.append(line)

        # 写入标签文件
        with open(txt_path, "w") as f:
            if lines:
                f.write("\n".join(lines) + "\n")

        print(f"✅ {img_filename} -> {len(lines)} 个检测框已保存")

    print(f"\n🎉 全部完成！标签文件保存在: {labels_dir}")


# ------------------- 主函数调用 -------------------
if __name__ == "__main__":
    MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_ailai_detect3/weights/best.pt"
    IMAGE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/ailaidete/train/delet"
    OUTPUT_DIR = "./inference_results"

    save_yolo_detect_labels_from_folder(
        model_path=MODEL_PATH,
        image_dir=IMAGE_DIR,
        output_dir=OUTPUT_DIR,
        conf_threshold=0.5
    )