# convert_dota_to_yolo_obb_manual.py
from pathlib import Path

def convert_dota_to_yolo_obb_manual(dota_root: str):
    dota_root = Path(dota_root)
    class_mapping = {"clamp": 0}  # 你可以加更多类别

    def convert_file(txt_path: Path, image_width: int, image_height: int, save_path: Path):
        with open(txt_path, 'r') as f:
            lines = f.readlines()

        with open(save_path, 'w') as g:
            for line in lines:
                parts = line.strip().split()
                if len(parts) < 9:
                    continue
                x1, y1, x2, y2, x3, y3, x4, y4 = map(float, parts[:8])
                class_name = parts[8]
                # difficult = parts[9]  # 可选，一般不用

                class_id = class_mapping.get(class_name.lower(), -1)
                if class_id == -1:
                    print(f"⚠️ 未知类别: {class_name}, 文件: {txt_path}")
                    continue

                # 归一化坐标
                coords = [x1, y1, x2, y2, x3, y3, x4, y4]
                normalized = [
                    coords[i] / (image_width if i % 2 == 0 else image_height)
                    for i in range(8)
                ]

                # 写入 YOLO-OBB 格式：class_id x1 y1 x2 y2 x3 y3 x4 y4
                line_str = f"{class_id} " + " ".join(f"{v:.6g}" for v in normalized)
                g.write(line_str + "\n")

    for phase in ["train", "val"]:
        image_dir = dota_root / "images" / phase
        label_dir = dota_root / "labels" / phase
        save_dir = dota_root / "labels_obb" / phase
        save_dir.mkdir(parents=True, exist_ok=True)

        for img_file in image_dir.iterdir():
            if img_file.suffix.lower() not in [".jpg", ".jpeg", ".png"]:
                continue

            # 获取图像宽高
            from PIL import Image
            img = Image.open(img_file)
            w, h = img.size

            # 找对应标签文件
            txt_file = label_dir / f"{img_file.stem}.txt"
            if not txt_file.exists():
                print(f"⚠️ 找不到标签文件: {txt_file}")
                continue

            save_file = save_dir / f"{img_file.stem}.txt"
            convert_file(txt_file, w, h, save_file)

    # 自动生成 dataset.yaml
    dataset_yaml = dota_root / "dataset.yaml"
    with open(dataset_yaml, 'w') as f:
        f.write(f"""train: {dota_root / 'images' / 'train'}\n""")
        f.write(f"""val: {dota_root / 'images' / 'val'}\n\n""")
        f.write(f"nc: {len(class_mapping)}\n")
        f.write(f"names: {list(class_mapping.keys())}\n")

    print(f"✅ 转换完成！标签已保存到: {dota_root / 'labels_obb'}")
    print(f"✅ dataset.yaml 已生成: {dataset_yaml}")

# 开始转换
convert_dota_to_yolo_obb_manual("/home/hx/桌面/image")