size_categories:n<1K format:parquet modality:image modality:text library:datasets library:pandas library:mlcroissant library:polars region:us
---
dataset_info:
  features:
  - name: image
    dtype: image
  - name: joy-caption
    dtype: string
  splits:
  - name: train
    num_bytes: 31576635.0
    num_examples: 106
  download_size: 31544968
  dataset_size: 31576635.0
configs:
- config_name: default
  data_files:
  - split: train
    path: data/train-*
---

```python
import os
import uuid
import toml
from PIL import Image
from datasets import load_dataset, Dataset
from tqdm import tqdm

# 1. 加载数据集
dataset_name = "svjack/Empress_Dowager_Cixi_Captioned"
ds = load_dataset(dataset_name)

# 2. 过滤数据集（如果需要）
#ds = ds.filter(lambda x: x["style_name"] in ["Comic book"])

# 3. 保存图片和文本文件
def save_image_and_text(dataset: Dataset, output_dir: str):
    """
    将数据集中的图片和文本保存为 PNG 和 TXT 文件。

    参数:
        dataset (Dataset): Hugging Face 数据集，包含 "image1", "image2" 和 "prompt" 列。
        output_dir (str): 输出文件的目录路径。
    """
    os.makedirs(output_dir, exist_ok=True)

    for example in tqdm(dataset):
        file_name_base = str(uuid.uuid4())

        # 保存 image1 和 prompt
        image1_path = os.path.join(output_dir, f"{file_name_base}_1.png")
        example["image"].save(image1_path)

        text1_path = os.path.join(output_dir, f"{file_name_base}_1.txt")
        with open(text1_path, "w", encoding="utf-8") as f:
            f.write(example["joy-caption"])

        print(f"Saved: {file_name_base}_1.png and {file_name_base}_1.txt")
        '''

        # 保存 image2 和 prompt
        image2_path = os.path.join(output_dir, f"{file_name_base}_2.png")
        example["image2"].save(image2_path)

        text2_path = os.path.join(output_dir, f"{file_name_base}_2.txt")
        with open(text2_path, "w", encoding="utf-8") as f:
            f.write(example["prompt"])

        print(f"Saved: {file_name_base}_2.png and {file_name_base}_2.txt")
        '''

# 4. 生成配置文件
def generate_image_config(image_dir: str, save_path: str = None):
    """
    生成图片配置文件的 TOML 格式。

    参数:
        image_dir (str): 图片目录路径。
        save_path (str): 配置文件的保存路径（可选）。
    """
    image_files = list(os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(".png"))
    if not image_files:
        raise ValueError("No PNG files found in the directory.")

    image_path = image_files[0]
    img = Image.open(image_path)
    width, height = img.size

    config = {
        "general": {
            "resolution": [width, height],
            "caption_extension": ".txt",
            "batch_size": 1,
            "enable_bucket": True,
            "bucket_no_upscale": False,
        },
        "datasets": [
            {
                "image_directory": image_dir,
            }
        ],
    }

    config_str = toml.dumps(config)
    print("Generated Configuration (TOML):")
    print(config_str)

    if save_path:
        with open(save_path, "w") as f:
            toml.dump(config, f)
        print(f"Configuration saved to {save_path}")

    img.close()

    return config_str

# 5. 根据数据集名称生成保存目录和配置文件名称
dataset_save_name = dataset_name.split("/")[-1]  # 提取数据集名称
output_dir = f"{dataset_save_name}_images_and_texts"
config_save_path = f"{dataset_save_name}_image_config.toml"

# 6. 保存图片和文本文件
save_image_and_text(ds["train"], output_dir)

# 7. 生成配置文件
generate_image_config(output_dir, config_save_path)

```
