#!/usr/bin/env python3
"""从 RF-DETR 候选标注整理 YOLO 训练集(class 0 确认野鸡 + class 1 疑似/藏身点, 80/20 拆分)"""
import os
import random
import shutil

BASE = os.path.dirname(os.path.abspath(__file__))
IMG_SRC = os.path.join(BASE, "datasets/images/pheasant")
LBL_SRC = os.path.join(BASE, "datasets/labels/pheasant_rfdetr")
OUT = os.path.join(BASE, "datasets/yolo")

random.seed(42)

stems = sorted(f[:-4] for f in os.listdir(LBL_SRC) if f.endswith(".txt"))
kept, dropped = [], []
for s in stems:
    src_lbl = os.path.join(LBL_SRC, s + ".txt")
    with open(src_lbl, encoding="utf-8") as f:
        lines = [l for l in f.read().splitlines() if l]
    if not lines:
        dropped.append(s)
        continue
    kept.append((s, lines))

random.shuffle(kept)
n_val = max(1, round(len(kept) * 0.2))
val, train = kept[:n_val], kept[n_val:]

for split, items in (("train", train), ("val", val)):
    for d in (os.path.join(OUT, "images", split), os.path.join(OUT, "labels", split)):
        os.makedirs(d, exist_ok=True)
    for s, lines in items:
        shutil.copy(os.path.join(IMG_SRC, s + ".jpg"), os.path.join(OUT, "images", split, s + ".jpg"))
        with open(os.path.join(OUT, "labels", split, s + ".txt"), "w", encoding="utf-8") as f:
            f.write("\n".join(lines) + "\n")

yaml = """path: /opt/pheasant_data/datasets/yolo
train: images/train
val: images/val
names:
  0: pheasant
  1: suspect
"""
with open(os.path.join(OUT, "dataset.yaml"), "w", encoding="utf-8") as f:
    f.write(yaml)

print(f"train: {len(train)} 张, val: {len(val)} 张, 空标注被剔除: {len(dropped)} 张")
print("剔除:", ", ".join(dropped) if dropped else "无")