Hugging Face's logo

Buckets:
twanghcmut
/
ego-real2sim

Files
xet
twanghcmut/ego-real2sim / basic_pick_place /code /examples /run_inverse.py
twanghcmut's picture
twanghcmut
download
raw
6.94 kB
import argparse
import json
import os
import subprocess
from datetime import datetime, timezone
from pathlib import Path
def main():
parser = argparse.ArgumentParser()
parser.add_argument("folder", type=Path)
parser.add_argument(
"--sim-python", type=Path, default=Path(".venv-mapanything/bin/python")
)
parser.add_argument(
"--sam-python", type=Path, default=Path(".venv-sam3/bin/python")
)
parser.add_argument(
"--sam-checkpoint", type=Path, default=Path("data/sam3/sam3.1_multiplex.pt")
)
parser.add_argument(
"--sam3d-root", type=Path, default=Path("external/sam-3d-objects")
)
parser.add_argument(
"--sam3d-checkpoints", type=Path, default=Path("data/sam3d/checkpoints")
)
parser.add_argument(
"--prompts", type=Path, default=Path("examples/basic_pick_place_prompts.json")
)
parser.add_argument("--device", default="cuda:1")
parser.add_argument("--population", type=int, default=16)
parser.add_argument("--iterations", type=int, default=6)
parser.add_argument("--frame", type=int)
parser.add_argument("--pixel", type=int, nargs=2)
parser.add_argument("--reuse", action="store_true")
parser.add_argument("--upload", action="store_true")
parser.add_argument("--token-file", type=Path)
args = parser.parse_args()
if (args.frame is None) != (args.pixel is None):
parser.error("frame and pixel must be supplied together")
folder = args.folder.resolve()
if not (folder / "sim_trajectory.npz").is_file():
parser.error("Run run_mapanything.py with hand tracks before inverse fitting")
environment = os.environ.copy()
environment["MUJOCO_GL"] = "egl"
environment["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
results = []
changed = set()
dependencies = {
"sam3d": {"sam31"},
"asset_preview": {"sam3d"},
"physics_search": {"sam31", "sam3d"},
"pixel_refinement": {"physics_search"},
"replay_verification": {"pixel_refinement"},
"pixel_viewer": {"pixel_refinement"},
"comparison_video": {"pixel_refinement"},
}
def run(stage, executable, script, arguments, artifact):
started = datetime.now(timezone.utc).isoformat()
reusable = (
args.reuse
and artifact.is_file()
and not dependencies.get(stage, set()).intersection(changed)
)
if reusable and stage == "pixel_refinement" and args.pixel is not None:
stored_query = json.loads(artifact.read_text()).get("pixel_query") or {}
reusable = (
stored_query.get("frame_index") == args.frame
and stored_query.get("pixel") == args.pixel
)
if reusable:
results.append(
{"stage": stage, "status": "reused", "artifact": str(artifact)}
)
else:
command = [
str(executable.absolute()),
str(Path(script).resolve()),
*map(str, arguments),
]
print(json.dumps({"stage": stage, "status": "running"}), flush=True)
subprocess.run(command, check=True, env=environment)
if not artifact.is_file():
raise RuntimeError(f"Missing output from {stage}: {artifact}")
changed.add(stage)
results.append(
{
"stage": stage,
"status": "complete",
"started_at": started,
"artifact": str(artifact),
}
)
(folder / "inverse_run.json").write_text(json.dumps(results, indent=2) + "\n")
run(
"sam31",
args.sam_python,
"examples/segment_mapanything.py",
[
folder,
"--prompts",
args.prompts,
"--checkpoint",
args.sam_checkpoint,
"--device",
args.device,
],
folder / "real_masks.npz",
)
run(
"sam3d",
args.sim_python,
"examples/reconstruct_object.py",
[
folder,
"--root",
args.sam3d_root,
"--checkpoints",
args.sam3d_checkpoints,
"--device",
args.device,
],
folder / "sam3d/object.npz",
)
run(
"asset_preview",
args.sim_python,
"examples/preview_object.py",
[folder / "sam3d/object.npz"],
folder / "sam3d/object_orbit.mp4",
)
run(
"physics_search",
args.sim_python,
"examples/fit_physics.py",
[
folder,
"--asset",
folder / "sam3d/object.npz",
"--population",
args.population,
"--iterations",
args.iterations,
],
folder / "physics/metrics.json",
)
refinement_arguments = [folder, "--asset", folder / "sam3d/object.npz"]
if args.pixel is not None:
refinement_arguments.extend(["--frame", args.frame, "--pixel", *args.pixel])
run(
"pixel_refinement",
args.sim_python,
"examples/refine_physics.py",
refinement_arguments,
folder / "physics/refinement/metrics.json",
)
run(
"replay_verification",
args.sim_python,
"examples/replay_physics.py",
[
folder,
"--asset",
folder / "sam3d/object.npz",
"--result",
folder / "physics/refinement",
],
folder / "physics/refinement/replay_verification.json",
)
run(
"pixel_viewer",
args.sim_python,
"examples/inspect_trace.py",
[
folder / "physics/refinement/refined_trace.npz",
"--video",
folder / "mujoco_replay.mp4",
"--metrics",
folder / "physics/refinement/metrics.json",
"--output",
folder / "physics/refinement/viewer.html",
],
folder / "physics/refinement/viewer.html",
)
run(
"comparison_video",
args.sim_python,
"examples/render_comparison.py",
[folder],
folder / "final_comparison.mp4",
)
if args.upload:
arguments = [
"--folder",
folder,
"--prefix",
"basic_pick_place",
"--sha256",
"--code",
]
if args.token_file:
arguments.extend(["--token-file", args.token_file])
subprocess.run(
[
str(args.sam_python.absolute()),
"examples/upload_results.py",
*map(str, arguments),
],
check=True,
env=environment,
)
print(json.dumps({"stages": results, "folder": str(folder)}), flush=True)
if __name__ == "__main__":
main()

Xet Storage Details

Size:
6.94 kB
·
Xet hash:
652b7b49516da3dcdc80f3317fadaca5b9423aa8a601073645c8f0d7861e3d44

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.