from datasets import load_dataset
import os

print("正在连接 Hugging Face 并提取高清数据集...")
dataset = load_dataset("Parveshiiii/AI-vs-Real", split="train", verification_mode="no_checks")

label_col = 'binary_label'
print(f"✅ 成功锁定标签列：'{label_col}'\n")

os.makedirs("./data/real_batch", exist_ok=True)
os.makedirs("./data/ai_batch", exist_ok=True)

# 设置我们想要的配额
target_count = 2000
real_saved = 0
ai_saved = 0

print(f"开始精准抓取：目标 {target_count}张真实图 + {target_count}张AI图...")

# 遍历整个数据集，直到两个配额都装满
for item in dataset:
    img = item['image']
    label = item[label_col]

    # 工业级防御：统一转为 RGB 格式
    if img.mode != 'RGB':
        img = img.convert('RGB')

    # 分流并计数
    if label == 1 and real_saved < target_count:
        img.save(f"./data/real_batch/real_hd_{real_saved}.jpg")
        real_saved += 1

    elif label == 0 and ai_saved < target_count:
        img.save(f"./data/ai_batch/ai_hd_{ai_saved}.jpg")
        ai_saved += 1

    # 每抓够 100 张，在终端汇报一下进度，让你心里有底
    if (real_saved + ai_saved) % 100 == 0:
        print(f"当前进度 -> 真实图: {real_saved}/{target_count} | AI图: {ai_saved}/{target_count}")

    # 如果两边的配额都满了，就停止遍历
    if real_saved >= target_count and ai_saved >= target_count:
        break

print("\n🎉 完美收工！这次两个文件夹绝对都塞得满满当当了！")