"""下载 TrashNet 数据集并整理为 ImageFolder 格式"""
from datasets import load_dataset
from pathlib import Path
import argparse

CLASS_NAMES = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]

def download_trashnet(output_dir="dataset/trashnet"):
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)

    print("正在从 Hugging Face 下载 TrashNet 数据集...")
    ds = load_dataset("garythung/trashnet", split="train", trust_remote_code=True)

    label_names = ds.features["label"].names
    print(f"类别: {label_names}")
    print(f"总样本数: {len(ds)}")

    # 为每个类别创建子目录
    for name in label_names:
        (out / name).mkdir(exist_ok=True)

    # 逐条保存
    for i, sample in enumerate(ds):
        label = label_names[sample["label"]]
        img = sample["image"]
        ext = "png" if img.mode == "RGBA" else "jpg"
        save_path = out / label / f"{label}_{i:05d}.{ext}"
        img = img.convert("RGB")
        img.save(save_path)
        if (i + 1) % 500 == 0:
            print(f"  已保存 {i + 1}/{len(ds)} 张...")

    # 打印统计
    print("\n下载完成！数据集统计:")
    for name in label_names:
        count = len(list((out / name).iterdir()))
        print(f"  {name}: {count} 张")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default="dataset/trashnet")
    args = parser.parse_args()
    download_trashnet(args.output_dir)