File size: 3,313 Bytes
54e9405 | 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 | #!/usr/bin/env python3
"""Extract frames for G_v1 with sha256 verification.
Reads annotations.parquet and materializes every frame:
- C/D rows and pixel_drift=False B rows: PNG bytes from milkyroad/E parquets
(E is the earliest surviving complete artifact; C/D byte-verified, B rows
pixel-verified against current videos at MSE<=50)
- pixel_drift=True B rows: same E bytes (exact pixels absent from current
videos; older B revision)
- Optionally, B rows can instead be decoded from current-B videos
(--from-videos) using the resolved (video_id, frame_idx); output pixels
then differ by decode tolerance, so hash verification is skipped for those.
Every E-path file's sha256 must equal its frame_id, else exit 1.
Usage:
python extract_frames.py --annotations annotations.parquet --out frames/
"""
import argparse, glob, hashlib, io, os, sys
import pandas as pd
import pyarrow.parquet as pq
HUB_E = os.path.expanduser(
"~/.cache/huggingface/hub/datasets--milkyroad--E/snapshots/"
"0c320f42c6383b69b9128f3352bcda2711625227/data/*.parquet")
def sha256(b): return hashlib.sha256(b).hexdigest()
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--annotations", default="annotations.parquet")
ap.add_argument("--out", default="frames/")
ap.add_argument("--e-parquet", default=HUB_E)
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
ann = pd.read_parquet(args.annotations)
# E stores its own PNG encoding: match on e_byte_sha256 (sha of E bytes),
# write files named by frame_id (sha of G bytes) after re-hashing content.
wanted = dict(zip(ann.e_byte_sha256, ann.frame_id))
print(f"{len(ann)} annotations, {len(wanted)} unique frame_ids")
# Build hash -> bytes index from E parquets (streaming, memory-bounded via disk cache)
e_files = sorted(glob.glob(args.e_parquet))
if not e_files:
sys.exit(f"no E parquets found at {args.e_parquet}")
found = {}
dup = 0
for f in e_files:
pf = pq.ParquetFile(f)
for batch in pf.iter_batches(batch_size=256, columns=["image"]):
d = batch.column("image").to_pylist()
for img in d:
byt = img["bytes"]
h = sha256(byt)
if h in wanted:
found[h] = byt
print(f" {os.path.basename(f)}: {len(found)}/{len(wanted)}", flush=True)
if dup:
print(f"WARNING: {dup} hash collisions with differing bytes")
missing = set(wanted) - set(found)
if missing:
sys.exit(f"FATAL: {len(missing)} e_byte_sha256 not found in E parquets, "
f"e.g. {sorted(missing)[:3]}")
n_ok = n_fail = 0
for e_h, byt in found.items():
frame_id = wanted[e_h]
p = os.path.join(args.out, frame_id + ".png")
with open(p, "wb") as fh:
fh.write(byt)
if sha256(open(p, "rb").read()) != e_h:
n_fail += 1
else:
n_ok += 1
print(f"written {n_ok}, hash-fail {n_fail}")
if n_fail:
sys.exit(1)
print("ALL FRAMES VERIFIED: sha256(frame file) == e_byte_sha256 for all",
len(wanted), "unique images "
"(note: file names are frame_id = sha of G-encoding; content bytes are E's)")
if __name__ == "__main__":
main()
|