File size: 3,824 Bytes
d0f5e0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | #!/usr/bin/env python3
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()
|