| |
| """Generate the 3000-video pendulum dataset + CLEVRER-style questions. |
| |
| Usage: python run_generate.py --root OUT --n-train 2400 --val-per-regime 100 --workers 2 |
| Resume: re-running the same command skips videos/CSVs that already exist. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import sys |
| import time |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| from pendulum_gen import (FRAMES, FPS, SIZE_PX, REGIMES_VAL, REGIME_DESCRIPTIONS, |
| generate_one) |
| from questions import build_questions |
|
|
| ROOT = None |
|
|
|
|
| def worker(job): |
| row, cfg, A = generate_one(job, str(ROOT)) |
| rng = np.random.default_rng(np.random.SeedSequence([job["seed"], job["index"], 777])) |
| qs = build_questions(cfg, A, rng) |
| entry = { |
| "video_id": job["video_id"], |
| "video_filename": row["video_path"], |
| "split": job["split"], |
| "regime": job["regime"], |
| "questions": qs, |
| } |
| return row, entry |
|
|
|
|
| def init_root(root): |
| global ROOT |
| ROOT = root |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--root", required=True) |
| ap.add_argument("--n-train", type=int, default=2400) |
| ap.add_argument("--val-per-regime", type=int, default=100) |
| ap.add_argument("--seed", type=int, default=20261014) |
| ap.add_argument("--workers", type=int, default=2) |
| args = ap.parse_args() |
|
|
| root = Path(args.root) |
| root.mkdir(parents=True, exist_ok=True) |
|
|
| jobs = [] |
| for i in range(args.n_train): |
| jobs.append({"seed": args.seed, "index": i, "split": "train", |
| "regime": "train", "video_id": f"video_{i:05d}"}) |
| k = 0 |
| for regime in REGIMES_VAL: |
| for j in range(args.val_per_regime): |
| jobs.append({"seed": args.seed, "index": 100000 + k, "split": "val", |
| "regime": regime, "video_id": f"video_{10000 + k:05d}"}) |
| k += 1 |
|
|
| print(f"{len(jobs)} videos -> {root}", flush=True) |
| t0 = time.time() |
| rows, q_entries = {}, {} |
| done = 0 |
| with Pool(args.workers, initializer=init_root, initargs=(root,)) as pool: |
| for row, entry in pool.imap_unordered(worker, jobs, chunksize=4): |
| rows[row["video_id"]] = row |
| q_entries[entry["video_id"]] = entry |
| done += 1 |
| if done % 25 == 0 or done == len(jobs): |
| rate = done / (time.time() - t0) |
| eta = (len(jobs) - done) / max(rate, 1e-9) |
| print(f"[{done}/{len(jobs)}] {rate:.2f} vid/s eta {eta/60:.1f} min", |
| flush=True) |
|
|
| |
| cols = ["video_id", "split", "regime", "video_path", "csv_path", "frames", |
| "duration_s", "fps", "size_px", "gravity_m_s2", "length_m", "name", |
| "theta0_deg", "omega0_rad_s", "damping_s_inv", "damping_model", |
| "quad_drag_beta", "drive_amp_rad_s2", "drive_freq_hz", "drive_phase_rad", |
| "light_mean_deg", "light_amplitude_deg", "light_frequency_hz", |
| "light_phase_rad"] |
| order = sorted(rows, key=lambda v: (0 if rows[v]["split"] == "train" else 1, v)) |
| with (root / "index.csv").open("w", newline="") as f: |
| w = csv.DictWriter(f, fieldnames=cols) |
| w.writeheader() |
| for v in order: |
| w.writerow(rows[v]) |
|
|
| |
| qroot = root / "questions" |
| qroot.mkdir(exist_ok=True) |
| for split in ("train", "val"): |
| entries = [q_entries[v] for v in order if q_entries[v]["split"] == split] |
| payload = { |
| "info": { |
| "description": "CLEVRER-style questions for the C-JEPA pendulum dataset", |
| "split": split, |
| "num_videos": len(entries), |
| "num_questions": sum(len(e["questions"]) for e in entries), |
| "question_types": ["descriptive", "explanatory", "predictive", |
| "counterfactual"], |
| "answer_convention": "descriptive: free-form string; others: index into 'choices'", |
| "answers_derived_from": "ground-truth RK4 simulation (28 s; videos show first 13.33 s)", |
| }, |
| "scene_annotation": entries, |
| } |
| with (qroot / f"questions_{split}.json").open("w") as f: |
| json.dump(payload, f, indent=1) |
|
|
| |
| meta = { |
| "dataset": "pendulum_cjepa_clevrer_3k", |
| "seed": args.seed, |
| "generation_mode": "random within per-regime bounds", |
| "num_videos": len(jobs), |
| "train_videos": args.n_train, |
| "val_videos": args.val_per_regime * len(REGIMES_VAL), |
| "val_design": "out-of-distribution physics regimes (NOT just a held-out seed)", |
| "val_regimes": REGIME_DESCRIPTIONS, |
| "frames_per_video": FRAMES, |
| "duration_s": FRAMES / FPS, |
| "fps": FPS, |
| "image_size": [SIZE_PX, SIZE_PX], |
| "codec": "H.264", "pixel_format": "yuv420p", "audio": False, |
| "ode": "theta_ddot = -(g/L) sin(theta) - gamma*theta_dot - beta*theta_dot*|theta_dot| + A*cos(2*pi*f_d*t+phi)", |
| "integrator": "RK4", "substeps_per_frame": 8, |
| "gravity_m_s2": "9.81 (train); {1.62, 3.71, 24.79} (val_gravity)", |
| "length_m": 1.0, |
| "split_unit": "whole video", |
| "normalized_labels": False, |
| "simulated_seconds_for_questions": 28.0, |
| "questions_per_video_target": 8, |
| } |
| with (root / "metadata.json").open("w") as f: |
| json.dump(meta, f, indent=2) |
|
|
| n_q = sum(len(e["questions"]) for e in q_entries.values()) |
| print(f"DONE {len(jobs)} videos, {n_q} questions, " |
| f"{(time.time()-t0)/60:.1f} min", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
