import os
import shutil
import random

# 输入路径
src_dir = r'F:\Cells_Identification\Selected_Data_Enhance'
dst_dir = r'F:\Cells_Identification\YOLO_ready'

# 支持的图片格式
img_exts = ['.jpg', '.jpeg', '.png', '.bmp']

# 创建目录结构
for sub in ['images/train', 'images/val', 'labels/train', 'labels/val']:
    os.makedirs(os.path.join(dst_dir, sub), exist_ok=True)

# 获取图片文件
all_imgs = [f for f in os.listdir(src_dir) if os.path.splitext(f)[1].lower() in img_exts]

# 洗牌并划分 8:2
random.shuffle(all_imgs)
split_idx = int(0.8 * len(all_imgs))
train_imgs = all_imgs[:split_idx]
val_imgs = all_imgs[split_idx:]

# 拷贝函数
def copy_files(file_list, split):
    for img_file in file_list:
        name, _ = os.path.splitext(img_file)
        label_file = name + '.txt'
        
        # 拷贝图片
        src_img = os.path.join(src_dir, img_file)
        dst_img = os.path.join(dst_dir, f'images/{split}', img_file)
        shutil.copy2(src_img, dst_img)

        # 拷贝标签（如果有）
        src_label = os.path.join(src_dir, label_file)
        dst_label = os.path.join(dst_dir, f'labels/{split}', label_file)
        if os.path.exists(src_label):
            shutil.copy2(src_label, dst_label)
        else:
            print(f'警告: 缺少标签 {label_file}，跳过。')

# 执行拷贝
copy_files(train_imgs, 'train')
copy_files(val_imgs, 'val')

# 读取 classes.txt 生成 cells.yaml
classes_path = os.path.join(src_dir, 'classes.txt')
if os.path.exists(classes_path):
    with open(classes_path, 'r', encoding='utf-8') as f:
        class_names = [line.strip() for line in f if line.strip()]
    nc = len(class_names)
else:
    class_names = []
    nc = 0
    print('警告: 未找到 classes.txt，将设置 nc=0，需手动编辑 cells.yaml')

yaml_path = os.path.join(dst_dir, 'cells.yaml')
with open(yaml_path, 'w', encoding='utf-8') as f:
    f.write(f"train: {dst_dir.replace(os.sep, '/')}/images/train\n")
    f.write(f"val: {dst_dir.replace(os.sep, '/')}/images/val\n\n")
    f.write(f"nc: {nc}\n")
    f.write(f"names: {class_names}\n")

print('✅ 数据整理完成，cells.yaml 已生成，可开始训练 YOLOv7。')