| |
| """Build WorldMem train/test CSV index files from mp4/npz pairs. |
| |
| Each row: video_path, action_path (paths relative to the project cwd when possible). |
| |
| When --root is a symlink under the project (e.g. ./data), CSV rows use paths like |
| data/training/.... Avoid resolving symlinks so relative paths stay valid. |
| |
| Example: |
| python scripts/build_worldmem_csv.py |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import random |
| from pathlib import Path |
|
|
|
|
| def _path_for_csv(path: Path) -> str: |
| """Path for CSV rows, relative to cwd when possible (preserves symlinks).""" |
| if not path.is_absolute(): |
| return path.as_posix() |
| try: |
| return path.relative_to(Path.cwd()).as_posix() |
| except ValueError: |
| return path.as_posix() |
|
|
|
|
| def collect_pairs(split_dir: Path, root: Path) -> list[tuple[str, str]]: |
| """Collect (video_path, action_path) for every mp4 with a matching npz.""" |
| if not split_dir.is_dir(): |
| raise FileNotFoundError(f"Split directory not found: {split_dir}") |
|
|
| pairs: list[tuple[str, str]] = [] |
| missing_npz: list[str] = [] |
|
|
| for mp4 in sorted(split_dir.rglob("*.mp4")): |
| npz = mp4.with_suffix(".npz") |
| if not npz.is_file(): |
| missing_npz.append(mp4.relative_to(root).as_posix()) |
| continue |
| pairs.append(( |
| _path_for_csv(mp4), |
| _path_for_csv(npz), |
| )) |
|
|
| if missing_npz: |
| print(f" warning: {len(missing_npz)} mp4 files without matching npz " |
| f"(first: {missing_npz[0]})") |
|
|
| return pairs |
|
|
|
|
| def write_csv(rows: list[tuple[str, str]], output_path: Path) -> None: |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| with output_path.open("w", newline="", encoding="utf-8") as f: |
| writer = csv.writer(f) |
| writer.writerow(["video_path", "action_path"]) |
| writer.writerows(rows) |
| print(f" wrote {len(rows)} rows -> {output_path}") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Build WorldMem train/test CSV metadata.") |
| parser.add_argument( |
| "--root", |
| type=str, |
| default='./data', |
| help=f"Dataset root directory", |
| ) |
| parser.add_argument( |
| "--output_dir", |
| type=str, |
| default='./data', |
| help="Directory for train.csv / test.csv (default: same as --root)", |
| ) |
| parser.add_argument( |
| "--seed", |
| type=int, |
| default=42, |
| help="Random seed for shuffling rows", |
| ) |
| args = parser.parse_args() |
|
|
| root = Path(args.root) |
| output_dir = Path(args.output_dir) if args.output_dir else root |
|
|
| splits = { |
| "train.csv": root / "training", |
| "test.csv": root / "test", |
| } |
|
|
| print(f"Dataset root : {root} (-> {root.resolve()})") |
| print(f"Output dir : {output_dir}") |
| print(f"Shuffle seed : {args.seed}") |
|
|
| rng = random.Random(args.seed) |
|
|
| for csv_name, split_dir in splits.items(): |
| print(f"\n[{csv_name}] scanning {split_dir.relative_to(root)}/") |
| pairs = collect_pairs(split_dir, root) |
| rng.shuffle(pairs) |
| write_csv(pairs, output_dir / csv_name) |
|
|
| print("\nDone.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
