#!/usr/bin/env python3
"""
Generate a JSON file listing exactly the SA-1B image paths
that have LayoutSAM annotations.
"""

import json
from datasets import load_dataset

def main():
    # 1. Load the LayoutSAM “train” split
    #    This gives us a column "image_path" with values like "sa_000000/sa_10000.jpg"
    ds = load_dataset("HuiZhang0812/LayoutSAM", split="train")

    # 2. Extract, dedupe, and sort the image_path entries
    image_paths = sorted({row["image_path"] for row in ds})

    # 3. Write out to JSON for the SA-1B-Downloader to consume
    out_file = "images_to_download.json"
    with open(out_file, "w") as f:
        json.dump(image_paths, f, indent=2)
    print(f"✅ Wrote {len(image_paths)} paths to {out_file}")

if __name__ == "__main__":
    main()