#!/usr/bin/env python3
"""
将红外图像数据集转换为YOLO格式
由于数据集没有标注，使用整个图像作为检测目标（目标检测模式）
"""

import argparse
import shutil
import random
from pathlib import Path
import cv2


def create_yolo_directories(output_dir: Path) -> dict:
    """创建YOLO目录结构"""
    dirs = {
        "train_images": output_dir / "images" / "train",
        "val_images": output_dir / "images" / "val",
        "train_labels": output_dir / "labels" / "train",
        "val_labels": output_dir / "labels" / "val",
    }
    for d in dirs.values():
        d.mkdir(parents=True, exist_ok=True)
    return dirs


def convert_dataset(input_dir: Path, output_dir: Path,
                   class_names: list,
                   train_ratio: float = 0.8,
                   use_whole_image: bool = True) -> None:
    """
    转换数据集为YOLO格式

    Args:
        input_dir: 输入图像目录（包含类别子目录）
        output_dir: 输出目录
        class_names: 类别名称列表
        train_ratio: 训练集比例
        use_whole_image: 是否使用整张图像作为检测目标
    """
    print("=" * 50)
    print("转换数据集为YOLO格式")
    print("=" * 50)
    print(f"输入目录: {input_dir}")
    print(f"输出目录: {output_dir}")
    print(f"使用整图作为目标: {use_whole_image}")
    print("=" * 50)

    # 创建输出目录
    dirs = create_yolo_directories(output_dir)

    # 创建类别到ID的映射
    class_to_id = {name: idx for idx, name in enumerate(class_names)}

    # 收集所有图像文件
    all_images = []
    for class_name in class_names:
        class_dir = input_dir / class_name
        if not class_dir.exists():
            print(f"警告: 目录不存在 {class_dir}")
            continue

        for img_file in class_dir.glob("*.jpg"):
            all_images.append((img_file, class_name))
        for img_file in class_dir.glob("*.png"):
            all_images.append((img_file, class_name))
        for img_file in class_dir.glob("*.jpeg"):
            all_images.append((img_file, class_name))

    print(f"找到 {len(all_images)} 张图像")

    if len(all_images) == 0:
        print("错误: 没有找到图像文件")
        return

    # 打乱并划分数据集
    random.seed(42)
    random.shuffle(all_images)

    train_count = int(len(all_images) * train_ratio)
    train_images = all_images[:train_count]
    val_images = all_images[train_count:]

    print(f"训练集: {len(train_images)} 张")
    print(f"验证集: {len(val_images)} 张")

    # 复制图像并创建标签
    for img_path, class_name in train_images:
        copy_and_label(img_path, class_name, dirs["train_images"],
                      dirs["train_labels"], use_whole_image, class_to_id)

    for img_path, class_name in val_images:
        copy_and_label(img_path, class_name, dirs["val_images"],
                      dirs["val_labels"], use_whole_image, class_to_id)

    # 创建dataset.yaml
    yaml_content = f"""# YOLO数据集配置
# 红外电力设备故障检测数据集

path: {output_dir.absolute()}
train: images/train
val: images/val

nc: {len(class_names)}
names: {class_names}
"""
    with open(output_dir / "dataset.yaml", "w", encoding="utf-8") as f:
        f.write(yaml_content)

    print("\n" + "=" * 50)
    print("转换完成!")
    print(f"配置文件: {output_dir / 'dataset.yaml'}")
    print("=" * 50)


def copy_and_label(img_path: Path, class_name: str,
                  output_img_dir: Path, output_label_dir: Path,
                  use_whole_image: bool,
                  class_to_id: dict) -> None:
    """复制图像并创建标签文件"""
    # 复制图像
    output_img_path = output_img_dir / img_path.name
    shutil.copy(img_path, output_img_path)

    # 创建标签文件
    label_path = output_label_dir / f"{img_path.stem}.txt"

    if use_whole_image:
        # 使用整张图像作为检测目标
        img = cv2.imread(str(img_path))
        if img is None:
            return

        h, w = img.shape[:2]

        # YOLO格式: class_id x_center y_center width height (归一化)
        # 整图作为目标，使用正确的类别ID
        class_id = class_to_id.get(class_name, 0)
        label_content = f"{class_id} 0.5 0.5 1.0 1.0\n"

        with open(label_path, "w") as f:
            f.write(label_content)
    else:
        # 空标签文件（用于后续添加真实标注）
        with open(label_path, "w") as f:
            pass


def main():
    parser = argparse.ArgumentParser(description="转换数据集为YOLO格式")
    parser.add_argument("--input-dir", type=str,
                       default="data/raw/images",
                       help="输入图像目录")
    parser.add_argument("--output-dir", type=str,
                       default="data/infrared",
                       help="输出目录")
    parser.add_argument("--train-ratio", type=float, default=0.8,
                       help="训练集比例")
    parser.add_argument("--whole-image", action="store_true",
                       help="使用整张图像作为检测目标")

    args = parser.parse_args()

    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir)

    if not input_dir.exists():
        print(f"错误: 输入目录不存在: {input_dir}")
        return

    # 从目录结构获取类别
    class_names = [d.name for d in input_dir.iterdir() if d.is_dir()]

    convert_dataset(input_dir, output_dir, class_names,
                  args.train_ratio, args.whole_image)


if __name__ == "__main__":
    main()