| |
| import argparse |
| import json |
| import os |
| import re |
| import time |
| from pathlib import Path |
|
|
| import pyarrow.parquet as pq |
|
|
|
|
| def safe_name(value: str, fallback: str) -> str: |
| name = Path(value or fallback).name or fallback |
| name = re.sub(r"[^A-Za-z0-9._-]+", "_", name) |
| return name or fallback |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--dataset-dir", required=True) |
| parser.add_argument("--image-root", required=True) |
| parser.add_argument("--json-out", required=True) |
| parser.add_argument("--batch-size", type=int, default=512) |
| parser.add_argument("--progress-every", type=int, default=5000) |
| parser.add_argument("--max-rows", type=int, default=0) |
| args = parser.parse_args() |
|
|
| dataset_dir = Path(args.dataset_dir).resolve() |
| image_root = Path(args.image_root).resolve() |
| json_out = Path(args.json_out).resolve() |
| tmp_out = json_out.with_suffix(json_out.suffix + ".tmp") |
|
|
| parquet_files = sorted(dataset_dir.rglob("*.parquet")) |
| if not parquet_files: |
| raise FileNotFoundError(f"No parquet files found under {dataset_dir}") |
|
|
| image_root.mkdir(parents=True, exist_ok=True) |
| json_out.parent.mkdir(parents=True, exist_ok=True) |
|
|
| total = 0 |
| started = time.time() |
| with tmp_out.open("w", encoding="utf-8") as fout: |
| fout.write("[\n") |
| first = True |
| for parquet_path in parquet_files: |
| part = parquet_path.parent.name |
| shard = parquet_path.stem |
| shard_dir = image_root / part / shard |
| shard_dir.mkdir(parents=True, exist_ok=True) |
|
|
| pf = pq.ParquetFile(parquet_path) |
| local_idx = 0 |
| for batch in pf.iter_batches( |
| batch_size=args.batch_size, |
| columns=["id", "conversations", "image_relpath", "image"], |
| ): |
| for row in batch.to_pylist(): |
| if args.max_rows and total >= args.max_rows: |
| break |
| relpath = row.get("image_relpath") or "" |
| basename = safe_name(relpath, f"{row.get('id') or total}.png") |
| image_path = shard_dir / f"{local_idx:06d}_{basename}" |
| image_obj = row.get("image") or {} |
| image_bytes = image_obj.get("bytes") |
| if image_bytes is None: |
| raise ValueError(f"Missing image bytes in {parquet_path} row {local_idx}") |
| if not image_path.exists() or image_path.stat().st_size != len(image_bytes): |
| image_path.write_bytes(image_bytes) |
|
|
| record = { |
| "id": row.get("id"), |
| "image": str(image_path), |
| "conversations": row.get("conversations"), |
| } |
| if first: |
| first = False |
| else: |
| fout.write(",\n") |
| json.dump(record, fout, ensure_ascii=False) |
|
|
| total += 1 |
| local_idx += 1 |
| if total % args.progress_every == 0: |
| elapsed = max(time.time() - started, 1e-6) |
| rate = total / elapsed |
| print(f"materialized_rows={total} rate={rate:.1f}/s elapsed={elapsed:.1f}s", flush=True) |
| if args.max_rows and total >= args.max_rows: |
| break |
| if args.max_rows and total >= args.max_rows: |
| break |
| fout.write("\n]\n") |
|
|
| os.replace(tmp_out, json_out) |
| elapsed = max(time.time() - started, 1e-6) |
| print(f"done rows={total} json={json_out} images={image_root} elapsed={elapsed:.1f}s", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
