#!/usr/bin/env python3
"""
为训练图片自动生成标注
使用 YOLO 模型检测野鸡并生成 YOLO 格式标注
"""

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

# 配置
MODEL_PATH = Path(__file__).parent / "yolov8s-world.pt"
IMAGES_DIR = Path(__file__).parent / "datasets" / "images" / "pheasant"
OUTPUT_DIR = Path(__file__).parent / "datasets" / "yolo_format"

def create_yolo_structure():
    """创建 YOLO 格式目录结构"""
    (OUTPUT_DIR / "images" / "train").mkdir(parents=True, exist_ok=True)
    (OUTPUT_DIR / "labels" / "train").mkdir(parents=True, exist_ok=True)

def convert_to_yolo_format(bbox, img_width, img_height):
    """将边界框转换为 YOLO 格式 (cx, cy, w, h)"""
    x1, y1, x2, y2 = bbox
    cx = (x1 + x2) / 2 / img_width
    cy = (y1 + y2) / 2 / img_height
    w = (x2 - x1) / img_width
    h = (y2 - y1) / img_height
    return cx, cy, w, h

def main():
    # 加载模型
    print(f"加载模型: {MODEL_PATH}")
    model = YOLO(str(MODEL_PATH))

    # 创建目录结构
    create_yolo_structure()

    # 获取所有图片
    image_extensions = {".jpg", ".jpeg", ".png", ".bmp"}
    image_files = [
        f for f in IMAGES_DIR.iterdir()
        if f.suffix.lower() in image_extensions
    ]

    print(f"找到 {len(image_files)} 张图片")

    # 处理每张图片
    labeled_count = 0
    for img_path in image_files:
        print(f"处理: {img_path.name}")

        # 读取图片
        image = cv2.imread(str(img_path))
        if image is None:
            print(f"  跳过: 无法读取 {img_path.name}")
            continue

        img_height, img_width = image.shape[:2]

        # 推理
        results = model(image, conf=0.3, iou=0.45)

        # 收集标注
        labels = []
        for result in results:
            boxes = result.boxes
            if boxes is None or len(boxes) == 0:
                continue

            for box in boxes:
                class_name = result.names.get(int(box.cls[0]), "")

                # 只保留 bird 或 pheasant 类别
                if class_name in ["bird", "pheasant"]:
                    class_id = 0  # 只有一个类别：野鸡
                    bbox = list(map(int, box.xyxy[0].tolist()))
                    conf = float(box.conf[0])

                    # 转换为 YOLO 格式
                    cx, cy, w, h = convert_to_yolo_format(bbox, img_width, img_height)
                    labels.append(f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")

        # 保存标注文件
        if len(labels) > 0:
            labeled_count += 1
            # 复制图片
            dst_img_path = OUTPUT_DIR / "images" / "train" / img_path.name
            cv2.imwrite(str(dst_img_path), image)

            # 保存标注
            label_path = OUTPUT_DIR / "labels" / "train" / (img_path.stem + ".txt")
            with open(label_path, "w") as f:
                f.write("\n".join(labels))

            print(f"  ✓ 标注了 {len(labels)} 个目标")
        else:
            print(f"  - 未检测到野鸡")

    print(f"\n完成!")
    print(f"  总图片数: {len(image_files)}")
    print(f"  有效标注: {labeled_count}")
    print(f"  输出目录: {OUTPUT_DIR}")

    # 创建 data.yaml
    data_yaml = OUTPUT_DIR / "data.yaml"
    with open(data_yaml, "w") as f:
        f.write(f"""# Observer 数据集配置 - 只识别野鸡
path: {OUTPUT_DIR}
train: images/train
val: images/train

# 类别
nc: 1
names: ['pheasant']
""")
    print(f"  数据配置: {data_yaml}")

if __name__ == "__main__":
    main()