Upload tools/sample_backgrounds.py with huggingface_hub

tools/sample_backgrounds.py

@@ -0,0 +1,918 @@
+import argparse
+import glob
+import json
+import logging
+import os
+import random
+import subprocess
+from io import BytesIO
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from PIL import Image
+
+from tools.dataset import BackgroundDataset, BackgroundIterableDataset
+
+
+logger = logging.getLogger(__name__)
+
+
+def iter_samples(dataset, streaming):
+    if streaming:
+        for sample in dataset:
+            yield sample
+    else:
+        for idx in range(len(dataset)):
+            yield dataset[idx]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Sample background images for SynLayers.")
+    parser.add_argument("--dataset-name", default="laion/laion2B-en-aesthetic")
+    parser.add_argument(
+        "--data-files",
+        default="/project/llmsvgen/share/data/kmw_layered_dataset/laion2B-en-aesthetic-image/*.parquet",
+        help="Parquet glob or list file.",
+    )
+    parser.add_argument("--split", default="train")
+    parser.add_argument("--cache-dir", default=None)
+    parser.add_argument("--url-column", default="URL")
+    parser.add_argument("--text-column", default="TEXT")
+    parser.add_argument("--hash-column", default="hash")
+    parser.add_argument(
+        "--image-root",
+        default="/project/llmsvgen/share/data/kmw_layered_dataset/laion2B-en-aesthetic-image",
+        help="Local directory with downloaded images named by hash.",
+    )
+    parser.add_argument(
+        "--image-extensions",
+        default=".jpg,.png,.jpeg,.webp",
+        help="Comma-separated extensions to try for local images.",
+    )
+    parser.add_argument("--image-size", type=int, default=None)
+    parser.add_argument("--count", type=int, default=10)
+    parser.add_argument("--streaming", action="store_true")
+    parser.add_argument("--output-dir", default="./outputs/backgrounds")
+    parser.add_argument(
+        "--save-images",
+        action="store_true",
+        help="Save images if found in image-root.",
+    )
+    parser.add_argument(
+        "--download",
+        action="store_true",
+        help="Download a subset into image-root using img2dataset.",
+    )
+    parser.add_argument(
+        "--download-mode",
+        choices=["auto", "img2dataset", "embedded"],
+        default="auto",
+        help="Download mode: auto-detect URL vs embedded bytes.",
+    )
+    parser.add_argument("--processes", type=int, default=8)
+    parser.add_argument("--threads", type=int, default=32)
+    parser.add_argument("--resize", type=int, default=512)
+    parser.add_argument("--build-splits", action="store_true")
+    parser.add_argument("--train-count", type=int, default=19000)
+    parser.add_argument("--val-count", type=int, default=1000)
+    parser.add_argument("--test-count", type=int, default=200)
+    parser.add_argument(
+        "--skip-existing",
+        action="store_true",
+        help="Skip downloading/extracting images that already exist in image-root.",
+    )
+    parser.add_argument(
+        "--progress-interval",
+        type=int,
+        default=500,
+        help="Log progress every N extracted images.",
+    )
+    parser.add_argument(
+        "--embedded-image-column",
+        default="whole_image",
+        help="Struct column containing embedded image bytes.",
+    )
+    parser.add_argument(
+        "--embedded-image-columns",
+        default=None,
+        help="Comma-separated embedded image columns to try in order.",
+    )
+    parser.add_argument(
+        "--embedded-image-bytes-key",
+        default="bytes",
+        help="Key inside embedded image struct that stores raw bytes.",
+    )
+    parser.add_argument(
+        "--embedded-image-path-key",
+        default="path",
+        help="Key inside embedded image struct that stores a path (if any).",
+    )
+    parser.add_argument(
+        "--embedded-caption-column",
+        default="whole_caption",
+        help="Caption column for embedded images.",
+    )
+    parser.add_argument(
+        "--embedded-id-column",
+        default="id",
+        help="ID column for embedded images.",
+    )
+    parser.add_argument(
+        "--size-multiple",
+        type=int,
+        default=8,
+        help="Round width/height up to a multiple of this value.",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--sequential",
+        action="store_true",
+        help="Use dataset order instead of random sampling when building splits.",
+    )
+    parser.add_argument(
+        "--allow-partial",
+        action="store_true",
+        help="Allow writing splits even if there are not enough images.",
+    )
+    parser.add_argument(
+        "--id-as-path",
+        action="store_true",
+        help="Store image path in the id field instead of the raw key.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
+    args = parse_args()
+
+    image_extensions = [ext.strip() for ext in args.image_extensions.split(",") if ext.strip()]
+
+    if args.download:
+        parquet_files = _expand_parquet_files(args.data_files)
+        if not parquet_files:
+            raise ValueError("No parquet files found. Check --data-files.")
+        os.makedirs(args.image_root, exist_ok=True)
+        download_mode = args.download_mode
+        if args.embedded_image_columns:
+            embedded_image_columns = [
+                col.strip() for col in args.embedded_image_columns.split(",") if col.strip()
+            ]
+        else:
+            embedded_image_columns = [args.embedded_image_column]
+        if download_mode == "auto":
+            if _parquet_has_column(parquet_files, args.url_column):
+                download_mode = "img2dataset"
+            elif any(
+                _parquet_has_column(parquet_files, col) for col in embedded_image_columns
+            ):
+                download_mode = "embedded"
+            else:
+                raise ValueError(
+                    "Could not detect download mode: missing URL and embedded image columns."
+                )
+        if download_mode == "img2dataset":
+            url_list_path = _prepare_download_parquet(
+                parquet_files=parquet_files,
+                output_dir=args.output_dir,
+                count=args.count,
+                seed=args.seed,
+                url_column=args.url_column,
+                text_column=args.text_column,
+                hash_column=args.hash_column,
+            )
+            cmd = [
+                "img2dataset",
+                "--url_list",
+                url_list_path,
+                "--input_format",
+                "parquet",
+                "--url_col",
+                args.url_column,
+                "--caption_col",
+                args.text_column,
+                "--output_format",
+                "files",
+                "--output_folder",
+                args.image_root,
+                "--processes_count",
+                str(args.processes),
+                "--thread_count",
+                str(args.threads),
+                "--image_size",
+                str(args.resize),
+                "--resize_mode",
+                "keep_ratio",
+            ]
+            logger.info("Downloading %d images into %s", args.count, args.image_root)
+            subprocess.run(cmd, check=True)
+        else:
+            logger.info(
+                "Extracting %d embedded images into %s",
+                args.count,
+                args.image_root,
+            )
+            download_embedded_images(
+                parquet_files=parquet_files,
+                image_root=args.image_root,
+                output_dir=args.output_dir,
+                count=args.count,
+                seed=args.seed,
+                sequential=args.sequential,
+                id_column=args.embedded_id_column,
+                caption_column=args.embedded_caption_column,
+                image_columns=embedded_image_columns,
+                image_bytes_key=args.embedded_image_bytes_key,
+                image_path_key=args.embedded_image_path_key,
+                image_extensions=image_extensions,
+                skip_existing=args.skip_existing,
+                progress_interval=args.progress_interval,
+            )
+
+    if args.build_splits:
+        if _has_img2dataset_parquet(args.image_root):
+            build_splits_from_img2dataset(
+                image_root=args.image_root,
+                output_dir=args.output_dir,
+                train_count=args.train_count,
+                val_count=args.val_count,
+                test_count=args.test_count,
+                seed=args.seed,
+                sequential=args.sequential,
+                allow_partial=args.allow_partial,
+                id_as_path=args.id_as_path,
+                image_extensions=image_extensions,
+                size_multiple=args.size_multiple,
+            )
+        else:
+            build_splits(
+                data_files=args.data_files,
+                image_root=args.image_root,
+                image_extensions=image_extensions,
+                output_dir=args.output_dir,
+                train_count=args.train_count,
+                val_count=args.val_count,
+                test_count=args.test_count,
+                seed=args.seed,
+                url_column=args.url_column,
+                text_column=args.text_column,
+                hash_column=args.hash_column,
+                sequential=args.sequential,
+                allow_partial=args.allow_partial,
+                size_multiple=args.size_multiple,
+            )
+        return
+
+    if args.streaming:
+        dataset = BackgroundIterableDataset(
+            dataset_name=args.dataset_name,
+            data_files=args.data_files,
+            split=args.split,
+            cache_dir=args.cache_dir,
+            url_column=args.url_column,
+            text_column=args.text_column,
+            hash_column=args.hash_column,
+            image_root=args.image_root,
+            image_extensions=image_extensions,
+            image_size=args.image_size,
+            require_image=args.save_images,
+        )
+    else:
+        dataset = BackgroundDataset(
+            dataset_name=args.dataset_name,
+            data_files=args.data_files,
+            split=args.split,
+            cache_dir=args.cache_dir,
+            url_column=args.url_column,
+            text_column=args.text_column,
+            hash_column=args.hash_column,
+            image_root=args.image_root,
+            image_extensions=image_extensions,
+            image_size=args.image_size,
+            max_items=args.count * 5,
+            require_image=args.save_images,
+        )
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    captions_path = os.path.join(args.output_dir, "captions.jsonl")
+
+    saved = 0
+    with open(captions_path, "w", encoding="utf-8") as captions_file:
+        for sample in iter_samples(dataset, args.streaming):
+            image = sample.get("image")
+            filename = None
+            if args.save_images:
+                if image is None:
+                    logger.warning("Skipping sample: local image not found.")
+                    continue
+                filename = f"background_{saved:03d}.png"
+                image.save(os.path.join(args.output_dir, filename))
+            captions_file.write(
+                json.dumps(
+                    {
+                        "file": filename,
+                        "url": sample.get("url"),
+                        "text": sample.get("text"),
+                        "width": sample.get("width"),
+                        "height": sample.get("height"),
+                        "hash": sample.get("hash"),
+                        "aesthetic": sample.get("aesthetic"),
+                        "punsafe": sample.get("punsafe"),
+                        "pwatermark": sample.get("pwatermark"),
+                    },
+                    ensure_ascii=False,
+                )
+                + "\n"
+            )
+            saved += 1
+            if saved >= args.count:
+                break
+
+    logger.info("Saved %d backgrounds to %s", saved, args.output_dir)
+
+
+def _expand_parquet_files(data_files):
+    if isinstance(data_files, (list, tuple)):
+        return list(data_files)
+    if not data_files:
+        return []
+    if os.path.exists(data_files) and data_files.endswith(".parquet"):
+        return [data_files]
+    return sorted(glob.glob(data_files))
+
+
+def _parquet_has_column(parquet_files, column_name):
+    if not column_name:
+        return False
+    for parquet_path in parquet_files:
+        parquet_file = pq.ParquetFile(parquet_path)
+        if column_name in parquet_file.schema.names:
+            return True
+        schema_arrow = getattr(parquet_file, "schema_arrow", None)
+        if schema_arrow is not None and column_name in schema_arrow.names:
+            return True
+    return False
+
+
+def _has_img2dataset_parquet(image_root):
+    if not image_root or not os.path.exists(image_root):
+        return False
+    return bool(glob.glob(os.path.join(image_root, "*.parquet")))
+
+
+def _prepare_download_parquet(
+    parquet_files,
+    output_dir,
+    count,
+    seed,
+    url_column,
+    text_column,
+    hash_column,
+):
+    os.makedirs(output_dir, exist_ok=True)
+    if len(parquet_files) == 1:
+        return parquet_files[0]
+    rng = random.Random(seed)
+    columns = [
+        url_column,
+        text_column,
+        hash_column,
+        "WIDTH",
+        "HEIGHT",
+        "aesthetic",
+        "punsafe",
+        "pwatermark",
+    ]
+    sampled = _reservoir_sample_parquet(
+        parquet_files=parquet_files,
+        target_count=count,
+        rng=rng,
+        columns=columns,
+    )
+    if not sampled:
+        raise ValueError("Failed to sample rows from parquet files.")
+    table = pa.Table.from_pylist(sampled)
+    out_path = os.path.join(output_dir, "laion_download_sample.parquet")
+    pq.write_table(table, out_path)
+    logger.info("Wrote sampled parquet list to %s", out_path)
+    return out_path
+
+
+def _detect_image_extension(image):
+    fmt = (image.format or "").upper()
+    if fmt == "JPEG":
+        return "jpg"
+    if fmt == "PNG":
+        return "png"
+    if fmt == "WEBP":
+        return "webp"
+    return "jpg"
+
+
+def _collect_existing_images(image_root, image_extensions):
+    if not image_root or not os.path.exists(image_root):
+        return {}
+    image_map = {}
+    for root, _, files in os.walk(image_root):
+        for name in files:
+            ext = os.path.splitext(name)[1].lower()
+            if ext in image_extensions:
+                stem = os.path.splitext(name)[0]
+                image_map[stem] = os.path.join(root, name)
+    return image_map
+
+
+def _save_image_bytes(image_bytes, output_path):
+    try:
+        with Image.open(BytesIO(image_bytes)) as img:
+            ext = _detect_image_extension(img)
+            if ext == "jpg":
+                img = img.convert("RGB")
+            elif img.mode not in ("RGB", "RGBA"):
+                img = img.convert("RGBA")
+            output_path = os.path.splitext(output_path)[0] + f".{ext}"
+            img.save(output_path)
+            return output_path, img.size
+    except Exception as exc:
+        logger.warning("Failed to decode image bytes: %s", exc)
+        return None, None
+
+
+def _iter_embedded_rows(
+    parquet_files,
+    id_column,
+    caption_column,
+    image_columns,
+    image_bytes_key,
+    image_path_key,
+):
+    columns = [id_column, caption_column] + list(image_columns)
+    for parquet_path in parquet_files:
+        parquet_file = pq.ParquetFile(parquet_path)
+        for batch in parquet_file.iter_batches(columns=columns, batch_size=256):
+            batch_dict = batch.to_pydict()
+            batch_len = len(batch)
+            for i in range(batch_len):
+                image_bytes = None
+                image_path = None
+                for image_column in image_columns:
+                    image_struct = batch_dict.get(image_column, [None])[i] or {}
+                    image_bytes = image_struct.get(image_bytes_key)
+                    image_path = image_struct.get(image_path_key)
+                    if image_bytes:
+                        break
+                if not image_bytes:
+                    continue
+                yield {
+                    "id": batch_dict.get(id_column, [None])[i],
+                    "caption": batch_dict.get(caption_column, [None])[i],
+                    "bytes": image_bytes,
+                    "path": image_path,
+                }
+
+
+def download_embedded_images(
+    parquet_files,
+    image_root,
+    output_dir,
+    count,
+    seed,
+    sequential,
+    id_column,
+    caption_column,
+    image_columns,
+    image_bytes_key,
+    image_path_key,
+    image_extensions,
+    skip_existing,
+    progress_interval,
+):
+    os.makedirs(image_root, exist_ok=True)
+    rng = random.Random(seed)
+    selected_ids = None
+    if not sequential:
+        sampled = _reservoir_sample_parquet(
+            parquet_files=parquet_files,
+            target_count=count,
+            rng=rng,
+            columns=[id_column],
+        )
+        selected_ids = {
+            str(row.get(id_column))
+            for row in sampled
+            if row.get(id_column) is not None
+        }
+        if not selected_ids:
+            raise ValueError("Failed to sample IDs from parquet files.")
+
+    image_extensions = image_extensions or [".jpg", ".png", ".jpeg", ".webp"]
+    existing_map = _collect_existing_images(image_root, image_extensions) if skip_existing else {}
+    if existing_map and len(existing_map) >= count:
+        logger.info(
+            "Found %d existing images in %s (target=%d).",
+            len(existing_map),
+            image_root,
+            count,
+        )
+    metadata_rows = []
+    for row in _iter_embedded_rows(
+        parquet_files=parquet_files,
+        id_column=id_column,
+        caption_column=caption_column,
+        image_columns=image_columns,
+        image_bytes_key=image_bytes_key,
+        image_path_key=image_path_key,
+    ):
+        image_id = row.get("id")
+        if image_id is None:
+            continue
+        image_id = str(image_id)
+        if selected_ids is not None and image_id not in selected_ids:
+            continue
+        saved_path = None
+        size = None
+        if image_id in existing_map:
+            saved_path = existing_map[image_id]
+            size = _get_image_size(saved_path)
+        if saved_path is None:
+            shard_dir = image_id[:5] if len(image_id) >= 5 else image_id
+            target_dir = os.path.join(image_root, shard_dir)
+            os.makedirs(target_dir, exist_ok=True)
+            target_path = os.path.join(target_dir, image_id)
+            saved_path, size = _save_image_bytes(row["bytes"], target_path)
+            if not saved_path:
+                continue
+        width, height = size if size else (None, None)
+        metadata_rows.append(
+            {
+                "key": image_id,
+                "caption": row.get("caption"),
+                "status": "success",
+                "width": width,
+                "height": height,
+            }
+        )
+        if progress_interval and len(metadata_rows) % progress_interval == 0:
+            logger.info("Extracted %d/%d images...", len(metadata_rows), count)
+        if sequential and len(metadata_rows) >= count:
+            break
+        if selected_ids is not None and len(metadata_rows) >= len(selected_ids):
+            break
+
+    if not metadata_rows:
+        raise ValueError("No embedded images were extracted.")
+    meta_table = pa.Table.from_pylist(metadata_rows)
+    meta_path = os.path.join(image_root, "embedded_metadata.parquet")
+    pq.write_table(meta_table, meta_path)
+    logger.info("Wrote embedded metadata to %s", meta_path)
+
+
+def _reservoir_sample_parquet(parquet_files, target_count, rng, columns):
+    sample = []
+    total_seen = 0
+    for parquet_path in parquet_files:
+        parquet_file = pq.ParquetFile(parquet_path)
+        for batch in parquet_file.iter_batches(columns=columns, batch_size=4096):
+            batch_dict = batch.to_pydict()
+            batch_len = len(batch)
+            for i in range(batch_len):
+                row = {col: batch_dict.get(col, [None])[i] for col in columns}
+                total_seen += 1
+                if len(sample) < target_count:
+                    sample.append(row)
+                else:
+                    j = rng.randint(0, total_seen - 1)
+                    if j < target_count:
+                        sample[j] = row
+    return sample
+
+
+def _iter_img2dataset_rows(image_root):
+    parquet_files = sorted(glob.glob(os.path.join(image_root, "*.parquet")))
+    if not parquet_files:
+        return
+    columns = ["key", "caption", "status", "width", "height"]
+    for parquet_path in parquet_files:
+        parquet_file = pq.ParquetFile(parquet_path)
+        for batch in parquet_file.iter_batches(columns=columns, batch_size=4096):
+            batch_dict = batch.to_pydict()
+            batch_len = len(batch)
+            for i in range(batch_len):
+                status = batch_dict.get("status", [None])[i]
+                if status and status != "success":
+                    continue
+                key = batch_dict.get("key", [None])[i]
+                caption = batch_dict.get("caption", [None])[i]
+                width = batch_dict.get("width", [None])[i]
+                height = batch_dict.get("height", [None])[i]
+                if key is None:
+                    continue
+                key_str = str(key)
+                yield {
+                    "id": key_str,
+                    "caption": caption,
+                    "width": width,
+                    "height": height,
+                }
+
+
+def _image_path_from_id(image_root, key_str, image_extensions):
+    if not key_str:
+        return None
+    shard_dir = key_str[:5]
+    for ext in image_extensions:
+        path = os.path.join(image_root, shard_dir, f"{key_str}{ext}")
+        if os.path.exists(path):
+            return path
+    return os.path.join(image_root, shard_dir, f"{key_str}.jpg")
+
+
+def _round_up_multiple(value, multiple):
+    if multiple <= 1:
+        return int(value)
+    return int(((value + multiple - 1) // multiple) * multiple)
+
+
+def _get_image_size(path):
+    try:
+        with Image.open(path) as img:
+            return img.size
+    except Exception as exc:
+        logger.warning("Failed to read image size for %s: %s", path, exc)
+        return None
+
+
+def build_splits_from_img2dataset(
+    image_root,
+    output_dir,
+    train_count,
+    val_count,
+    test_count,
+    seed,
+    sequential=False,
+    allow_partial=False,
+    id_as_path=False,
+    image_extensions=None,
+    size_multiple=8,
+):
+    os.makedirs(output_dir, exist_ok=True)
+    total_needed = train_count + val_count + test_count
+    image_extensions = image_extensions or [".jpg", ".png", ".jpeg", ".webp"]
+    items = []
+    if sequential:
+        for row in _iter_img2dataset_rows(image_root):
+            items.append(row)
+            if len(items) >= total_needed:
+                break
+    else:
+        rng = random.Random(seed)
+        total_seen = 0
+        for row in _iter_img2dataset_rows(image_root):
+            total_seen += 1
+            if len(items) < total_needed:
+                items.append(row)
+            else:
+                j = rng.randint(0, total_seen - 1)
+                if j < total_needed:
+                    items[j] = row
+        rng.shuffle(items)
+
+    if len(items) < total_needed:
+        if not allow_partial:
+            raise ValueError(
+                f"Only found {len(items)} matching images (needed {total_needed})."
+            )
+        logger.warning(
+            "Only found %d matching images (needed %d).",
+            len(items),
+            total_needed,
+        )
+
+    if id_as_path:
+        for item in items:
+            item["id"] = _image_path_from_id(image_root, item["id"], image_extensions)
+
+    train_items = items[:train_count]
+    val_items = items[train_count : train_count + val_count]
+    test_items = items[train_count + val_count : train_count + val_count + test_count]
+
+    def write_jsonl(path, rows):
+        with open(path, "w", encoding="utf-8") as f:
+            for row in rows:
+                image_path = row.get("path")
+                if not image_path:
+                    image_id = row.get("id")
+                    if image_id:
+                        if os.path.isabs(image_id):
+                            image_path = image_id
+                        else:
+                            image_path = _image_path_from_id(
+                                image_root, image_id, image_extensions
+                            )
+                if image_path:
+                    row["path"] = image_path
+                    size = _get_image_size(image_path)
+                    if size:
+                        width, height = size
+                    else:
+                        width = row.get("width")
+                        height = row.get("height")
+                    if width and height:
+                        row["width"] = _round_up_multiple(int(width), size_multiple)
+                        row["height"] = _round_up_multiple(int(height), size_multiple)
+                f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+    write_jsonl(os.path.join(output_dir, "train.jsonl"), train_items)
+    write_jsonl(os.path.join(output_dir, "val.jsonl"), val_items)
+    write_jsonl(os.path.join(output_dir, "test.jsonl"), test_items)
+
+    logger.info(
+        "Wrote splits to %s (train=%d, val=%d, test=%d)",
+        output_dir,
+        len(train_items),
+        len(val_items),
+        len(test_items),
+    )
+
+
+def _scan_images(image_root, image_extensions):
+    if not image_root or not os.path.exists(image_root):
+        return {}
+    image_map = {}
+    for root, _, files in os.walk(image_root):
+        for name in files:
+            ext = os.path.splitext(name)[1].lower()
+            if ext in image_extensions:
+                stem = os.path.splitext(name)[0]
+                image_map[stem] = os.path.join(root, name)
+    return image_map
+
+
+def _collect_metadata(
+    parquet_files,
+    image_map,
+    target_count,
+    url_column,
+    text_column,
+    hash_column,
+):
+    selected = []
+    hashes = set(image_map.keys())
+    if not hashes:
+        return selected
+    columns = [
+        hash_column,
+        url_column,
+        text_column,
+        "WIDTH",
+        "HEIGHT",
+        "aesthetic",
+        "punsafe",
+        "pwatermark",
+    ]
+    for parquet_path in parquet_files:
+        parquet_file = pq.ParquetFile(parquet_path)
+        for batch in parquet_file.iter_batches(columns=columns, batch_size=4096):
+            batch_dict = batch.to_pydict()
+            for i in range(len(batch)):
+                hash_value = batch_dict.get(hash_column, [None])[i]
+                if hash_value is None:
+                    continue
+                hash_str = str(hash_value)
+                path = image_map.get(hash_str)
+                if not path:
+                    continue
+                selected.append(
+                    {
+                        "file": path,
+                        "url": batch_dict.get(url_column, [None])[i],
+                        "text": batch_dict.get(text_column, [None])[i],
+                        "width": batch_dict.get("WIDTH", [None])[i],
+                        "height": batch_dict.get("HEIGHT", [None])[i],
+                        "hash": hash_str,
+                        "aesthetic": batch_dict.get("aesthetic", [None])[i],
+                        "punsafe": batch_dict.get("punsafe", [None])[i],
+                        "pwatermark": batch_dict.get("pwatermark", [None])[i],
+                    }
+                )
+                if len(selected) >= target_count:
+                    return selected
+    return selected
+
+
+def build_splits(
+    data_files,
+    image_root,
+    image_extensions,
+    output_dir,
+    train_count,
+    val_count,
+    test_count,
+    seed,
+    url_column,
+    text_column,
+    hash_column,
+    sequential=False,
+    allow_partial=False,
+    size_multiple=8,
+):
+    os.makedirs(output_dir, exist_ok=True)
+    parquet_files = _expand_parquet_files(data_files)
+    if not parquet_files:
+        raise ValueError("No parquet files found. Check --data-files.")
+
+    image_map = _scan_images(image_root, image_extensions)
+    if not image_map:
+        raise ValueError("No images found in image_root.")
+
+    total_needed = train_count + val_count + test_count
+    logger.info(
+        "Collecting %d samples from %d parquet files (images=%d)",
+        total_needed,
+        len(parquet_files),
+        len(image_map),
+    )
+    items = _collect_metadata(
+        parquet_files=parquet_files,
+        image_map=image_map,
+        target_count=total_needed,
+        url_column=url_column,
+        text_column=text_column,
+        hash_column=hash_column,
+    )
+    if len(items) < total_needed:
+        if not allow_partial:
+            raise ValueError(
+                f"Only found {len(items)} matching images (needed {total_needed})."
+            )
+        logger.warning(
+            "Only found %d matching images (needed %d).",
+            len(items),
+            total_needed,
+        )
+
+    if not sequential:
+        rng = random.Random(seed)
+        rng.shuffle(items)
+    train_items = items[:train_count]
+    val_items = items[train_count : train_count + val_count]
+    test_items = items[train_count + val_count : train_count + val_count + test_count]
+
+    def write_jsonl(path, rows):
+        with open(path, "w", encoding="utf-8") as f:
+            for row in rows:
+                image_path = row.get("path") or row.get("file")
+                if image_path:
+                    row["path"] = image_path
+                    size = _get_image_size(image_path)
+                    if size:
+                        width, height = size
+                    else:
+                        width = row.get("width")
+                        height = row.get("height")
+                    if width and height:
+                        row["width"] = _round_up_multiple(int(width), size_multiple)
+                        row["height"] = _round_up_multiple(int(height), size_multiple)
+                f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+    write_jsonl(os.path.join(output_dir, "train.jsonl"), train_items)
+    write_jsonl(os.path.join(output_dir, "val.jsonl"), val_items)
+    write_jsonl(os.path.join(output_dir, "test.jsonl"), test_items)
+
+    logger.info(
+        "Wrote splits to %s (train=%d, val=%d, test=%d)",
+        output_dir,
+        len(train_items),
+        len(val_items),
+        len(test_items),
+    )
+
+
+if __name__ == "__main__":
+    main()
+
+'''
+python -m tools.sample_backgrounds \
+  --download \
+  --count 20100 \
+  --build-splits \
+  --train-count 19000 \
+  --val-count 1000 \
+  --test-count 200 \
+  --data-files "/project/llmsvgen/share/data/kmw_layered_dataset/laion2B-en-aesthetic-image/*.parquet" \
+  --image-root "/project/llmsvgen/share/data/kmw_layered_dataset/laion2B-en-aesthetic-image" \
+  --output-dir "/project/llmsvgen/jinmin/SynLayers/data/laion2b_splits"
+
+python -m tools.sample_backgrounds \
+  --download \
+  --build-splits \
+  --count 40200 \
+  --sequential \
+  --id-as-path \
+  --train-count 19000 \
+  --val-count 1000 \
+  --test-count 200 \
+  --data-files "/project/llmsvgen/share/data/kmw_layered_dataset/PrismLayersPro-image/data/*.parquet" \
+  --image-root "/project/llmsvgen/share/data/kmw_layered_dataset/PrismLayersPro-image/data/haolin/PrismLayersPro-image" \
+  --output-dir "/project/llmsvgen/jinmin/SynLayers/data/prismlayerspro_splits"
+'''