49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
#!/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 "无")
|