Function main

motion_gating/render_mbench_videos.py:128–171

()

def main() -> None:
    args = parse_args()
    meta_json = args.meta_json.expanduser().resolve()
    output_json = args.output_json.expanduser().resolve()
    output_dir = args.output_dir.expanduser().resolve()
    smplx_model_dir = args.smplx_model_dir.expanduser().resolve()
    repo_root = find_repo_root(meta_json)

    entries = json.loads(meta_json.read_text())
    if not isinstance(entries, list):
        raise ValueError(f"Expected a list in {meta_json}, got {type(entries).__name__}")

    updated_entries: List[Dict[str, Any]] = []
    to_process = entries if args.limit is None else entries[: args.limit]

    for entry in tqdm(to_process, desc="Render MBench videos"):
        motion_path = entry.get("motion_path")
        if not motion_path:
            updated_entries.append(entry)
            continue
        motion_file = resolve_repo_path(meta_json, str(motion_path))
        # Use global_id to ensure unique file names across dimensions
        motion_id = entry.get("global_id", entry.get("id", Path(motion_file).stem))
        outputs = render_single_motion(
            motion_file,
            motion_id,
            output_dir,
            smplx_model_dir,
            args.device,
            fps=args.fps,
            width=args.width,
            height=args.height,
        )
        updated_entry = update_entry_paths(entry, repo_root, Path(outputs["pt"]), Path(outputs["mp4"]))
        updated_entries.append(updated_entry)

    # Preserve any remaining entries if limit was used
    if args.limit is not None and len(entries) > args.limit:
        updated_entries.extend(entries[args.limit :])

    output_json.parent.mkdir(parents=True, exist_ok=True)
    output_json.write_text(json.dumps(updated_entries, indent=2))
    print(f"Rendered videos/PT caches to {output_dir}")
    print(f"Wrote updated meta JSON with video_path/mbench_eval_path to {output_json}")


if __name__ == "__main__":