Spaces:
Paused
Paused
File size: 2,312 Bytes
20857b0 | 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 | """Export .gtkv token streams to Parquet / Hugging Face datasets.
The Parquet schema mirrors the container's hierarchical token layout —
one row per chunk, one column per token layer. AI models can consume
the dataset directly without touching pixels.
"""
from pathlib import Path
from mediatok.container.gtkv import GtkvReader
from mediatok.entropy import entropy_decode
LAYER_COLUMNS = ["scene", "camera", "object", "motion", "texture", "residual"]
def gtkv_to_rows(gtkv_path: str) -> list[dict]:
"""Read a .gtkv file and return one dict per chunk."""
reader = GtkvReader(gtkv_path)
h = reader.header
rows = []
try:
for ci in range(reader.num_chunks):
block = reader.read_video_block(ci)
tokens = entropy_decode(block.entropy_payload, h.entropy_codec_id,
block.token_count)
row = {
"chunk_index": ci,
"num_frames": h.chunk_size_frames,
"width": h.width,
"height": h.height,
"fps": h.fps,
}
offset = 0
for li, col in enumerate(LAYER_COLUMNS):
n = block.layer_sizes[li] if li < len(block.layer_sizes) else 0
row[f"{col}_tokens"] = tokens[offset:offset + n] if n > 0 else []
offset += n
rows.append(row)
finally:
reader.close()
return rows
def gtkv_to_parquet(gtkv_path: str, out_path: str) -> str:
"""Write a .gtkv token stream to a Parquet file."""
import pyarrow as pa
import pyarrow.parquet as pq
rows = gtkv_to_rows(gtkv_path)
table = pa.Table.from_pylist(rows)
pq.write_table(table, out_path)
return out_path
def gtkv_to_hf_dataset(gtkv_path: str):
"""Return a Hugging Face Dataset of the token stream."""
from datasets import Dataset
return Dataset.from_list(gtkv_to_rows(gtkv_path))
def push_to_hub(gtkv_path: str, repo_id: str, token: str | None = None,
private: bool = False):
"""Upload the token dataset to the Hugging Face Hub (Parquet-backed).
Requires `huggingface_hub` credentials (HF_TOKEN env var or token arg).
"""
ds = gtkv_to_hf_dataset(gtkv_path)
return ds.push_to_hub(repo_id, token=token, private=private)
|