diff --git a/README.md b/README.md index 8c168d9e..08c16e31 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ work can use `skip_vertices=True` to avoid preparing rest vertices. - Anatomicals: SKEL, MyoFullBody - Heads: FLAME - Hands: MANO -- Robots: BrainCo, G1 +- Robots: BrainCo, G1, SmplHumanoid See the [model docs](https://abcamiletto.github.io/body-models/#supported-models) for setup, supported backends, inputs, and model-specific behavior. diff --git a/docs/index.md b/docs/index.md index 22137a33..499eb435 100644 --- a/docs/index.md +++ b/docs/index.md @@ -56,6 +56,7 @@ uv add "body-models[jax]" | --- | --- | --- | | [BrainCo](models/brainco.md) | BrainCo Revo 2 robotic hand | auto-download | | [G1](models/g1.md) | Unitree G1 rigid links | auto-download | +| [SmplHumanoid](models/smpl-humanoid.md) | SMPL-compatible humanoid MJCF variants | auto-download | ## Common Usage diff --git a/docs/models/anny.md b/docs/models/anny.md index 36407cc5..1ca62b63 100644 --- a/docs/models/anny.md +++ b/docs/models/anny.md @@ -4,7 +4,9 @@ ANNY is a phenotype-driven body model with configurable rig and topology variant ## Setup -ANNY downloads automatically on first use. To prefetch and save the path: +ANNY downloads automatically on first use from `https://huggingface.co/abcamiletto/body-models`. +The Hugging Face repo records the original ANNY Apache 2.0 and MPFB2 CC0 provenance. To prefetch +and save the path: ```bash # Download the ANNY assets and store their path in the body-models config. diff --git a/docs/models/brainco.md b/docs/models/brainco.md index 6b28e3af..6b9e9088 100644 --- a/docs/models/brainco.md +++ b/docs/models/brainco.md @@ -4,7 +4,7 @@ BrainCo is a rigid articulated model of the BrainCo Revo 2 robotic hand using th ## Setup -BrainCo downloads automatically on first use. To prefetch and save the path: +BrainCo downloads from the public [`abcamiletto/body-models`](https://huggingface.co/abcamiletto/body-models) Hugging Face repository on first use. To prefetch and save the path: ```bash # Download the BrainCo MuJoCo XML and STL assets. @@ -13,6 +13,8 @@ body-models download brainco When passed manually, `model_path` should contain `left.xml`, `right.xml`, and `meshes/{left,right}/*.STL`. +The original BrainCo Revo2 description license is included with the hosted assets. + ## Usage ```python diff --git a/docs/models/g1.md b/docs/models/g1.md index 07299a51..57270ceb 100644 --- a/docs/models/g1.md +++ b/docs/models/g1.md @@ -4,7 +4,9 @@ G1 is a rigid articulated Unitree G1 model with STL link meshes attached to the ## Setup -G1 downloads automatically on first use. To prefetch and save the path: +G1 downloads automatically on first use from `https://huggingface.co/abcamiletto/body-models`. +The Hugging Face repo records the original GR00T-WholeBodyControl / LeRobot provenance. To prefetch +and save the path: ```bash # Download the Unitree G1 XML and link meshes. diff --git a/docs/models/garment-measurements.md b/docs/models/garment-measurements.md index 461acb73..c9887e5d 100644 --- a/docs/models/garment-measurements.md +++ b/docs/models/garment-measurements.md @@ -4,8 +4,9 @@ GarmentMeasurements is a PCA body model with an FBX-derived skeleton and skinnin ## Setup -GarmentMeasurements downloads its preprocessed asset from -`https://huggingface.co/datasets/abcamiletto/body-models-assets` on first use. To prefetch and save the path: +GarmentMeasurements downloads its preprocessed assets from +`https://huggingface.co/abcamiletto/body-models` on first use. The Hugging Face repo records +the original SOMA-X Apache 2.0 provenance for the source asset. To prefetch and save the path: ```bash # Download the preprocessed GarmentMeasurements body-model asset. diff --git a/docs/models/mhr.md b/docs/models/mhr.md index 02937c03..d30dd6cd 100644 --- a/docs/models/mhr.md +++ b/docs/models/mhr.md @@ -4,13 +4,17 @@ MHR is an expressive full-body model with neural pose correctives. ## Setup -MHR downloads automatically on first use. To prefetch and save the path: +MHR downloads from the public [`abcamiletto/body-models`](https://huggingface.co/abcamiletto/body-models) Hugging Face repository on first use. The hosted package keeps the original MHR checkpoint for the default LOD 1 path and adds preprocessed FBX-derived mesh assets for LODs 0 through 6. + +To prefetch and save the path: ```bash # Download the MHR assets and store their path in the body-models config. body-models download mhr ``` +The original MHR license is included with the hosted assets. + ## API ::: body_models.bodies.mhr.numpy.MHR diff --git a/docs/models/myofullbody.md b/docs/models/myofullbody.md index 977d47fb..81011bf8 100644 --- a/docs/models/myofullbody.md +++ b/docs/models/myofullbody.md @@ -4,7 +4,8 @@ MyoFullBody is a MuJoCo-derived musculoskeletal full-body model from `amathislab ## Setup -MyoFullBody downloads automatically on first use. To prefetch and save the path: +MyoFullBody downloads automatically on first use from `https://huggingface.co/abcamiletto/body-models`. +The Hugging Face repo records the original MuscleMimic Apache 2.0 provenance. To prefetch and save the path: ```bash # Download the MyoFullBody MJCF and referenced mesh assets. diff --git a/docs/models/smpl-humanoid.md b/docs/models/smpl-humanoid.md new file mode 100644 index 00000000..037879ed --- /dev/null +++ b/docs/models/smpl-humanoid.md @@ -0,0 +1,18 @@ +# SmplHumanoid + +SmplHumanoid is a rigid articulated humanoid model loaded from SMPL-compatible MJCF XML variants. + +## Setup + +SmplHumanoid downloads its XML assets from the public [`abcamiletto/body-models`](https://huggingface.co/abcamiletto/body-models) Hugging Face repository. To prefetch and save the hosted path: + +```bash +# Download the SmplHumanoid MJCF XML assets. +body-models download smpl-humanoid +``` + +The hosted folder includes license/provenance notes for the XML variants. + +## API + +::: body_models.robots.smpl_humanoid.numpy.SmplHumanoid diff --git a/docs/models/soma.md b/docs/models/soma.md index 76b1fd8d..b9fc9647 100644 --- a/docs/models/soma.md +++ b/docs/models/soma.md @@ -4,7 +4,8 @@ SOMA provides a native implementation for SOMA-X assets with identity, pose, and ## Setup -SOMA downloads automatically on first use. To prefetch and save the path: +SOMA downloads automatically on first use from `https://huggingface.co/abcamiletto/body-models`. +The Hugging Face repo records the original SOMA-X Apache 2.0 provenance. To prefetch and save the path: ```bash # Download the SOMA-X assets used by the native SOMA implementation. diff --git a/src/body_models/__init__.py b/src/body_models/__init__.py index 40bdf791..bf62d46a 100644 --- a/src/body_models/__init__.py +++ b/src/body_models/__init__.py @@ -67,6 +67,7 @@ def main() -> None: from .config import CONFIG_FILE, MODELS, get_model_path, set_model_path, unset_model_path from .robots.brainco.io import download_model as download_brainco_model from .robots.g1.io import download_model as download_g1_model + from .robots.smpl_humanoid.io import download_model as download_smpl_humanoid_model from .skeletons.myofullbody.io import download_model as download_myofullbody_model Model = Literal[ @@ -88,6 +89,7 @@ def main() -> None: "flame", "brainco", "g1", + "smpl-humanoid", "soma", "garment-measurements", "myofullbody", @@ -130,6 +132,7 @@ def download( "brainco", "mhr", "g1", + "smpl-humanoid", "soma", "garment-measurements", "myofullbody", @@ -234,6 +237,11 @@ def download( set_model_path("g1", str(path)) print(f"Set g1 = {path}") + if model in ("smpl-humanoid", "all"): + path = download_smpl_humanoid_model() + set_model_path("smpl-humanoid", str(path)) + print(f"Set smpl-humanoid = {path}") + if model in ("soma", "all"): path = download_soma_model() set_model_path("soma", str(path)) diff --git a/src/body_models/bodies/anny/io.py b/src/body_models/bodies/anny/io.py index 217d108b..97f1948e 100644 --- a/src/body_models/bodies/anny/io.py +++ b/src/body_models/bodies/anny/io.py @@ -11,14 +11,14 @@ from nanomanifold import SO3 from body_models import config -from body_models.cache import download_and_extract, get_cache_dir +from body_models.cache import HF_MODEL_BASE_URL, download_and_extract, get_cache_dir PathLike = Path | str Array = Any Front = tuple[list[int], list[int]] # One FK depth level: (joint_indices, parent_indices). -ANNY_URL = "https://github.com/naver/anny/archive/refs/heads/main.zip" +ANNY_URL = f"{HF_MODEL_BASE_URL}/anny/assets.zip" PHENOTYPE_VARIATIONS = { "race": ["african", "asian", "caucasian"], @@ -83,7 +83,7 @@ def get_model_path(model_path: PathLike | None = None) -> Path: def download_model() -> Path: cache_dir = get_cache_dir() / "anny" print(f"Downloading ANNY model to {cache_dir}...") - download_and_extract(url=ANNY_URL, dest=cache_dir, extract_subdir="anny-main/src/anny/") + download_and_extract(url=ANNY_URL, dest=cache_dir) print("Done") return cache_dir diff --git a/src/body_models/bodies/garment_measurements/io.py b/src/body_models/bodies/garment_measurements/io.py index 99dabc05..75bb7981 100644 --- a/src/body_models/bodies/garment_measurements/io.py +++ b/src/body_models/bodies/garment_measurements/io.py @@ -10,7 +10,7 @@ from jaxtyping import Float, Int from body_models import config -from body_models.cache import HF_DATASET_BASE_URL, download_file, get_cache_dir +from body_models.cache import HF_MODEL_BASE_URL, download_and_extract, get_cache_dir PathLike = Path | str Array = Any @@ -18,7 +18,7 @@ Front = tuple[list[int], list[int]] # One FK depth level: (joint_indices, parent_indices). PREPROCESSED_FILENAME = "garment_measurements.npz" -GARMENT_MEASUREMENTS_URL = f"{HF_DATASET_BASE_URL}/garment_measurements/{PREPROCESSED_FILENAME}" +GARMENT_MEASUREMENTS_URL = f"{HF_MODEL_BASE_URL}/garment_measurements/assets.zip" @dataclass(frozen=True) @@ -65,7 +65,7 @@ def download_model() -> Path: """Download preprocessed GarmentMeasurements data assets.""" cache_dir = get_cache_dir() / "garment_measurements" print(f"Downloading GarmentMeasurements model to {cache_dir}...") - download_file(GARMENT_MEASUREMENTS_URL, cache_dir / PREPROCESSED_FILENAME) + download_and_extract(GARMENT_MEASUREMENTS_URL, cache_dir) print("Done") return cache_dir diff --git a/src/body_models/bodies/mhr/io.py b/src/body_models/bodies/mhr/io.py index b7c3cda9..711d5185 100644 --- a/src/body_models/bodies/mhr/io.py +++ b/src/body_models/bodies/mhr/io.py @@ -11,7 +11,7 @@ from body_models import config from body_models.common import simplify_mesh -from body_models.cache import download_and_extract, get_cache_dir +from body_models.cache import HF_MODEL_BASE_URL, download_and_extract, get_cache_dir PathLike = Path | str @@ -28,7 +28,8 @@ "load_pose_correctives", ] -MHR_URL = "https://github.com/facebookresearch/MHR/releases/download/v1.0.0/assets.zip" +MHR_URL = f"{HF_MODEL_BASE_URL}/mhr/assets.zip" +SUPPORTED_LODS = tuple(range(7)) @dataclass(frozen=True) @@ -71,17 +72,17 @@ def get_model_path(model_path: PathLike | None = None) -> Path: return validate_path(model_path) cache_path = get_cache_dir() / "mhr" - if (cache_path / "mhr_model.pt").exists(): + if _has_hosted_assets(cache_path): return cache_path return download_model() def download_model() -> Path: - """Download MHR model from GitHub releases.""" + """Download MHR model assets.""" cache_dir = get_cache_dir() / "mhr" print(f"Downloading MHR model to {cache_dir}...") - download_and_extract(url=MHR_URL, dest=cache_dir, extract_subdir="assets/") + download_and_extract(url=MHR_URL, dest=cache_dir) print("Done") return cache_dir @@ -89,10 +90,11 @@ def download_model() -> Path: def load_model_data(asset_dir: Path, *, lod: int = 1, simplify: float = 1.0) -> MhrWeights: if simplify < 1.0: raise ValueError("simplify must be >= 1.0") - if lod != 1: - raise ValueError("MHR lod values other than 1 are not supported.") + if lod not in SUPPORTED_LODS: + raise ValueError(f"MHR lod must be one of {SUPPORTED_LODS}, got {lod}") - data = _load_raw_model_data(asset_dir) + shared_data = _load_raw_model_data(asset_dir) + data = shared_data if lod == 1 else _load_preprocessed_lod_data(asset_dir, lod, shared_data) base_vertices = data["base_vertices"] blendshape_dirs = data["blendshape_dirs"] skin_weights = data["skin_weights"] @@ -100,6 +102,10 @@ def load_model_data(asset_dir: Path, *, lod: int = 1, simplify: float = 1.0) -> faces = data["faces"].astype(np.int64) corrective_weights = load_pose_correctives_weights(asset_dir, lod) corrective_W2 = corrective_weights["W2"] + if corrective_W2.shape[0] != len(base_vertices) * 3: + raise ValueError( + f"MHR lod{lod} corrective W2 has {corrective_W2.shape[0]} rows, expected {len(base_vertices) * 3}" + ) if simplify > 1.0: target_faces = int(len(faces) / simplify) @@ -169,6 +175,43 @@ def _load_raw_model_data(asset_dir: Path) -> dict[str, Any]: } +def _load_preprocessed_lod_data(asset_dir: Path, lod: int, shared_data: dict[str, Any]) -> dict[str, Any]: + path = asset_dir / f"mhr_lod{lod}.npz" + with np.load(path, allow_pickle=False) as asset: + joint_names = [str(name) for name in asset["skin_joint_names"].tolist()] + checkpoint_joint_index = {name: index for index, name in enumerate(shared_data["joint_names"])} + missing = sorted(name for name in joint_names if name not in checkpoint_joint_index) + if missing: + raise ValueError(f"{path} references joints missing from mhr_model.pt: {missing}") + + skin_joint_indices = np.asarray(asset["skin_joint_indices"], dtype=np.int64) + mapped_joint_indices = np.asarray([checkpoint_joint_index[joint_names[index]] for index in skin_joint_indices]) + base_vertices = np.asarray(asset["base_vertices"], dtype=np.float32) + skin_indices, skin_weights = _build_dense_skinning( + asset["skin_vertex_indices"], + mapped_joint_indices, + asset["skin_weights"], + len(base_vertices), + ) + + return shared_data | { + "base_vertices": base_vertices, + "blendshape_dirs": np.asarray(asset["blendshape_dirs"], dtype=np.float32), + "skin_weights": skin_weights, + "skin_indices": skin_indices, + "faces": np.asarray(asset["faces"], dtype=np.int64), + } + + +def _has_hosted_assets(model_path: Path) -> bool: + asset_names = [ + "mhr_model.pt", + "corrective_activation.npz", + *(f"mhr_lod{lod}.npz" for lod in SUPPORTED_LODS), + ] + return all((model_path / name).is_file() for name in asset_names) + + def _get_attr(obj: Any, path: str) -> Any: cur = obj for part in path.split("."): diff --git a/src/body_models/bodies/soma/io.py b/src/body_models/bodies/soma/io.py index f24f95ea..d74a8728 100644 --- a/src/body_models/bodies/soma/io.py +++ b/src/body_models/bodies/soma/io.py @@ -18,7 +18,7 @@ from body_models import config from body_models.common import simplify_mesh -from body_models.cache import HF_DATASET_BASE_URL, download_file, get_cache_dir +from body_models.cache import HF_MODEL_BASE_URL, download_and_extract, get_cache_dir PathLike = Path | str @@ -30,7 +30,7 @@ SOMA_PROCEDURAL_TRANSFORMS_ASSET = "SOMA_procedural_transforms.json" SOMA_ASSETS = (SOMA_CORE_ASSET, SOMA_CORRECTIVES_ASSET) SOMA_UPSTREAM_02_ASSETS = (SOMA_TEMPLATE_RIG_ASSET, SOMA_PROCEDURAL_TRANSFORMS_ASSET) -SOMA_BASE_URL = f"{HF_DATASET_BASE_URL}/soma" +SOMA_URL = f"{HF_MODEL_BASE_URL}/soma/assets.zip" SOMA_LEGACY_NPZ_FIELDS = ( "bind_shape", "bind_pose_world", @@ -274,12 +274,9 @@ def get_model_path(model_path: PathLike | None = None) -> Path: def download_model(model_dir: PathLike | None = None) -> Path: """Download SOMA assets from Hugging Face.""" cache_dir = Path(model_dir) if model_dir is not None else get_cache_dir() / "soma" - cache_dir.mkdir(parents=True, exist_ok=True) - missing = [name for name in SOMA_ASSETS if not (cache_dir / name).exists()] - if missing: + if _missing_assets(cache_dir): print(f"Downloading SOMA model to {cache_dir}...") - for name in missing: - download_file(f"{SOMA_BASE_URL}/{name}", cache_dir / name) + download_and_extract(url=SOMA_URL, dest=cache_dir) print("Done") return validate_path(cache_dir) @@ -288,21 +285,15 @@ def ensure_identity_assets(model_dir: Path, model_type: str) -> None: """Ensure supplementary SOMA assets exist for a given identity backend.""" normalized = model_type.lower() spec = MODEL_TYPE_SPECS.get(normalized) - if spec is None or spec.asset_dir is None: + if spec is None or spec.asset_dir is None or spec.source_mesh_name is None or spec.target_mesh_name is None: raise ValueError(f"Unsupported SOMA identity assets: {model_type}") asset_dir = Path(model_dir) - asset_names = ( - f"{spec.asset_dir}/{spec.target_mesh_name}", - f"{spec.asset_dir}/{spec.source_mesh_name}", - ) - missing = [name for name in asset_names if not (asset_dir / name).exists()] - if missing: - print(f"Downloading SOMA {normalized} assets to {asset_dir}...") - for name in missing: - path = asset_dir / name - download_file(f"{SOMA_BASE_URL}/{name}", path) - print("Done") + identity_dir = asset_dir / spec.asset_dir + source_mesh = identity_dir / spec.source_mesh_name + target_mesh = identity_dir / spec.target_mesh_name + if not source_mesh.exists() or not target_mesh.exists(): + download_model(asset_dir) def preprocess_model(upstream_dir: PathLike, output_dir: PathLike) -> Path: @@ -416,16 +407,7 @@ def _procedural_rig_data(data: Any) -> SomaProceduralRig | None: "twist_axis_signs": np.float32, } values = {name: np.asarray(data[f"procedural_{name}"], dtype=dtypes[name]) for name in SOMA_PROCEDURAL_RIG_FIELDS} - return SomaProceduralRig( - public_joint_indices_full=values["public_joint_indices_full"], - rotation_matrix=values["rotation_matrix"], - translation_matrix=values["translation_matrix"], - source_axis_ids=values["source_axis_ids"], - source_axis_signs=values["source_axis_signs"], - twist_joint_indices=values["twist_joint_indices"], - twist_axis_ids=values["twist_axis_ids"], - twist_axis_signs=values["twist_axis_signs"], - ) + return SomaProceduralRig(**values) def with_active_mesh( diff --git a/src/body_models/cache.py b/src/body_models/cache.py index f27bb0e4..73a255e0 100644 --- a/src/body_models/cache.py +++ b/src/body_models/cache.py @@ -2,9 +2,15 @@ from platformdirs import user_cache_dir -__all__ = ["get_cache_dir", "get_cached_path", "download_file", "download_and_extract"] +__all__ = [ + "HF_MODEL_BASE_URL", + "get_cache_dir", + "get_cached_path", + "download_file", + "download_and_extract", +] -HF_DATASET_BASE_URL = "https://huggingface.co/datasets/abcamiletto/body-models-assets/resolve/main/models_hub" +HF_MODEL_BASE_URL = "https://huggingface.co/abcamiletto/body-models/resolve/main" def get_cache_dir() -> Path: @@ -34,19 +40,8 @@ def download_file(url: str, dest: Path) -> None: shutil.copyfileobj(src, dst) -def download_and_extract( - url: str, - dest: Path, - extract_subdir: str | None = None, -) -> None: - """Download a zip file and extract it to dest. - - Args: - url: URL to download from. - dest: Destination directory for extracted files. - extract_subdir: If specified, only extract files from this subdirectory - within the zip archive. The subdirectory prefix is stripped. - """ +def download_and_extract(url: str, dest: Path) -> None: + """Download a zip file and extract it to dest.""" import tempfile import zipfile @@ -59,27 +54,6 @@ def download_and_extract( download_file(url, tmp_path) with zipfile.ZipFile(tmp_path) as zf: - if extract_subdir is None: - zf.extractall(dest) - else: - # Normalize the subdir path - if not extract_subdir.endswith("/"): - extract_subdir = extract_subdir + "/" - - for member in zf.namelist(): - if member.startswith(extract_subdir): - # Strip the subdirectory prefix - relative_path = member[len(extract_subdir) :] - if not relative_path: - continue - - target_path = dest / relative_path - - if member.endswith("/"): - target_path.mkdir(parents=True, exist_ok=True) - else: - target_path.parent.mkdir(parents=True, exist_ok=True) - with zf.open(member) as src, open(target_path, "wb") as dst: - dst.write(src.read()) + zf.extractall(dest) finally: tmp_path.unlink(missing_ok=True) diff --git a/src/body_models/config.py b/src/body_models/config.py index 9aa92628..cc594522 100644 --- a/src/body_models/config.py +++ b/src/body_models/config.py @@ -26,6 +26,7 @@ "flame", "brainco", "g1", + "smpl-humanoid", "soma", "garment-measurements", "myofullbody", @@ -86,6 +87,8 @@ def validate_model_path(model: str, path: str | Path) -> Path: from .robots.brainco.io import validate_path elif model == "g1": from .robots.g1.io import validate_path + elif model == "smpl-humanoid": + from .robots.smpl_humanoid.io import validate_path elif model == "soma": from .bodies.soma.io import validate_path elif model == "garment-measurements": diff --git a/src/body_models/robots/brainco/io.py b/src/body_models/robots/brainco/io.py index c6321d33..209f477a 100644 --- a/src/body_models/robots/brainco/io.py +++ b/src/body_models/robots/brainco/io.py @@ -3,10 +3,7 @@ from __future__ import annotations from dataclasses import dataclass -import shutil -import urllib.request import xml.etree.ElementTree as ET -import zipfile from pathlib import Path from typing import Any, Literal @@ -14,7 +11,7 @@ from jaxtyping import Float, Int from body_models import config -from body_models.cache import get_cache_dir +from body_models.cache import HF_MODEL_BASE_URL, download_and_extract, get_cache_dir from body_models.common.stl import load_stl_mesh as _load_stl_mesh from body_models.robots import mjcf @@ -22,9 +19,7 @@ Side = Literal["left", "right"] Array = Any VALID_SIDES = ("left", "right") -BRAINCO_REVO2_MUJOCO_URL = ( - "https://brainco-common-public.oss-cn-hangzhou.aliyuncs.com/web-config/docs-sdk/Revo2_xml.zip" -) +BRAINCO_URL = f"{HF_MODEL_BASE_URL}/brainco/assets.zip" MUJOCO_TO_KIMODO = np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float32) JOINT_SUFFIXES = [ "base_skel", @@ -91,31 +86,11 @@ def get_model_path(model_path: PathLike | None = None) -> Path: def download_model() -> Path: - """Download the official BrainCo Revo 2 MuJoCo model package.""" + """Download the BrainCo Revo 2 MuJoCo model package.""" cache_dir = get_cache_dir() / "brainco" - if cache_dir.exists(): - shutil.rmtree(cache_dir) - cache_dir.mkdir(parents=True) - archive_path = cache_dir / "revo2_mujoco.zip" - urllib.request.urlretrieve(BRAINCO_REVO2_MUJOCO_URL, archive_path) - with zipfile.ZipFile(archive_path) as zf: - for member in zf.infolist(): - if member.is_dir(): - continue - side = _archive_side(member.filename) - if side is None: - continue - name = Path(member.filename).name - if name.endswith(".xml") and name.startswith("brainco-"): - target = cache_dir / f"{side}.xml" - elif "/meshes/" in member.filename and name.endswith(".STL"): - target = cache_dir / "meshes" / side / _mesh_name(side, name) - else: - continue - target.parent.mkdir(parents=True, exist_ok=True) - with zf.open(member) as src, target.open("wb") as dst: - shutil.copyfileobj(src, dst) - archive_path.unlink(missing_ok=True) + print(f"Downloading BrainCo model to {cache_dir}...") + download_and_extract(url=BRAINCO_URL, dest=cache_dir) + print("Done") return validate_path(cache_dir) @@ -398,14 +373,6 @@ def _joint_limit(joint: ET.Element, class_limits: dict[str, tuple[float, float]] raise ValueError(f"Missing limit for BrainCo joint {joint.get('name')}") -def _archive_side(filename: str) -> str | None: - if "/xml_left/" in filename: - return "left" - if "/xml_right/" in filename: - return "right" - return None - - def _mesh_name(side: str, name: str) -> str: return name if name.startswith(f"{side}_") else f"{side}_{name}" diff --git a/src/body_models/robots/g1/io.py b/src/body_models/robots/g1/io.py index c341be5d..d9e0b955 100644 --- a/src/body_models/robots/g1/io.py +++ b/src/body_models/robots/g1/io.py @@ -3,7 +3,6 @@ from __future__ import annotations from dataclasses import dataclass -import urllib.request import xml.etree.ElementTree as ET from pathlib import Path from typing import Any, Literal @@ -12,7 +11,7 @@ from jaxtyping import Float, Int from body_models import config -from body_models.cache import get_cache_dir +from body_models.cache import HF_MODEL_BASE_URL, download_and_extract, get_cache_dir from body_models.common.stl import load_stl_mesh as _load_stl_mesh from body_models.robots import mjcf @@ -22,8 +21,7 @@ MUJOCO_TO_KIMODO = np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float32) VALID_CONVENTIONS = ("soma", "mujoco") -G1_HF_BASE_URL = "https://huggingface.co/lerobot/unitree-g1-mujoco/resolve/main/assets" -G1_HF_XML = "g1_29dof_no_hand.xml" +G1_URL = f"{HF_MODEL_BASE_URL}/g1/assets.zip" JOINT_NAMES = [ "pelvis_skel", @@ -172,14 +170,8 @@ def get_model_path(model_path: PathLike | None = None) -> Path: def download_model() -> Path: """Download G1 XML and STL assets from Hugging Face.""" cache_dir = get_cache_dir() / "g1" - mesh_dir = cache_dir / "meshes" - cache_dir.mkdir(parents=True, exist_ok=True) - mesh_dir.mkdir(parents=True, exist_ok=True) - print(f"Downloading G1 model to {cache_dir}...") - urllib.request.urlretrieve(f"{G1_HF_BASE_URL}/{G1_HF_XML}", cache_dir / "g1.xml") - for mesh_name in sorted({mesh for meshes in G1_MESH_JOINT_MAP.values() for mesh in meshes}): - urllib.request.urlretrieve(f"{G1_HF_BASE_URL}/meshes/{mesh_name}", mesh_dir / mesh_name) + download_and_extract(url=G1_URL, dest=cache_dir) print("Done") return cache_dir / "g1.xml" diff --git a/src/body_models/robots/smpl_humanoid/__init__.py b/src/body_models/robots/smpl_humanoid/__init__.py index e5d68e42..3897b9af 100644 --- a/src/body_models/robots/smpl_humanoid/__init__.py +++ b/src/body_models/robots/smpl_humanoid/__init__.py @@ -2,13 +2,11 @@ from body_models.robots.smpl_humanoid.constants import SMPL_HUMANOID_VARIANTS from body_models.robots.smpl_humanoid.io import ( - SMPL_HUMANOID_SOURCES, SmplHumanoidWeights, load_model_data, ) __all__ = [ - "SMPL_HUMANOID_SOURCES", "SMPL_HUMANOID_VARIANTS", "SmplHumanoidWeights", "load_model_data", diff --git a/src/body_models/robots/smpl_humanoid/assets/xml/humenv.xml b/src/body_models/robots/smpl_humanoid/assets/xml/humenv.xml deleted file mode 100644 index 9e7d1f20..00000000 --- a/src/body_models/robots/smpl_humanoid/assets/xml/humenv.xml +++ /dev/null @@ -1,316 +0,0 @@ - - diff --git a/src/body_models/robots/smpl_humanoid/assets/xml/phc.xml b/src/body_models/robots/smpl_humanoid/assets/xml/phc.xml deleted file mode 100644 index eb42a587..00000000 --- a/src/body_models/robots/smpl_humanoid/assets/xml/phc.xml +++ /dev/null @@ -1,245 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/src/body_models/robots/smpl_humanoid/assets/xml/smplsim.xml b/src/body_models/robots/smpl_humanoid/assets/xml/smplsim.xml deleted file mode 100644 index 49a0cee9..00000000 --- a/src/body_models/robots/smpl_humanoid/assets/xml/smplsim.xml +++ /dev/null @@ -1,294 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/src/body_models/robots/smpl_humanoid/io.py b/src/body_models/robots/smpl_humanoid/io.py index f1fa8c0e..9068a728 100644 --- a/src/body_models/robots/smpl_humanoid/io.py +++ b/src/body_models/robots/smpl_humanoid/io.py @@ -12,13 +12,14 @@ import trimesh.creation from trimesh import Trimesh +from body_models import config +from body_models.cache import HF_MODEL_BASE_URL, download_and_extract, get_cache_dir from body_models.robots import mjcf from body_models.robots.smpl_humanoid.constants import BODY_JOINTS, JOINT_NAMES, PARENTS, SMPL_HUMANOID_VARIANTS Array = Any PathLike = Path | str -XML_DIR = Path(__file__).parent / "assets" / "xml" -SMPL_HUMANOID_SOURCES: dict[str, Path] = {name: XML_DIR / f"{name}.xml" for name in SMPL_HUMANOID_VARIANTS} +SMPL_HUMANOID_URL = f"{HF_MODEL_BASE_URL}/smpl_humanoid/assets.zip" @dataclass(frozen=True) @@ -43,6 +44,18 @@ class SmplHumanoidWeights: actuated_joint_types: list[str] +@dataclass(frozen=True) +class LinkData: + joint_indices: list[int] + vertex_starts: list[int] + vertex_counts: list[int] + face_starts: list[int] + face_counts: list[int] + geom_positions: Float[Array, "L 3"] + geom_rotations: Float[Array, "L 3 3"] + names: list[str] + + def load_model_data(source: PathLike = "humenv", *, dtype=np.float32) -> SmplHumanoidWeights: """Load a rigid SMPL humanoid from an MJCF XML file.""" path = _model_source(source) @@ -76,7 +89,7 @@ def load_model_data(source: PathLike = "humenv", *, dtype=np.float32) -> SmplHum if parsed_parent_indices != PARENTS: raise ValueError("SMPL humanoid XML body hierarchy does not match the canonical SMPL hierarchy.") - vertices, faces, link_data = _load_xml_geoms(parsed_bodies, dtype=dtype) + vertices, faces, links = _load_xml_geoms(parsed_bodies, dtype=dtype) actuated_joint_indices = [by_name[name] for name, _ in BODY_JOINTS] actuated_joint_names = [name for name, _ in BODY_JOINTS for _ in range(3)] num_actuated = 3 * len(BODY_JOINTS) @@ -90,14 +103,14 @@ def load_model_data(source: PathLike = "humenv", *, dtype=np.float32) -> SmplHum rest_local_rotations=rest_local_rotations.astype(dtype), vertices=vertices.astype(dtype), faces=faces.astype(np.int64), - link_joint_indices=link_data["joint_indices"], - link_vertex_starts=link_data["vertex_starts"], - link_vertex_counts=link_data["vertex_counts"], - link_face_starts=link_data["face_starts"], - link_face_counts=link_data["face_counts"], - link_geom_positions=link_data["geom_positions"].astype(dtype), - link_geom_rotations=link_data["geom_rotations"].astype(dtype), - link_names=link_data["names"], + link_joint_indices=links.joint_indices, + link_vertex_starts=links.vertex_starts, + link_vertex_counts=links.vertex_counts, + link_face_starts=links.face_starts, + link_face_counts=links.face_counts, + link_geom_positions=links.geom_positions.astype(dtype), + link_geom_rotations=links.geom_rotations.astype(dtype), + link_names=links.names, actuated_joint_indices=actuated_joint_indices, actuated_joint_limits=actuated_joint_limits, actuated_joint_names=actuated_joint_names, @@ -105,19 +118,58 @@ def load_model_data(source: PathLike = "humenv", *, dtype=np.float32) -> SmplHum ) +def validate_path(path: PathLike) -> Path: + path = Path(path) + if path.is_file(): + raise ValueError(f"Expected an SMPL humanoid XML directory, got file: {path}") + if not path.is_dir(): + raise FileNotFoundError(f"SMPL humanoid model path not found: {path}") + missing = [f"{name}.xml" for name in SMPL_HUMANOID_VARIANTS if not (path / f"{name}.xml").is_file()] + if missing: + raise FileNotFoundError(f"SMPL humanoid model path {path} is missing: {', '.join(missing)}") + return path + + +def get_model_path(model_path: PathLike | None = None) -> Path: + if model_path is None: + model_path = config.get_model_path("smpl-humanoid") + if model_path is not None: + return validate_path(model_path) + + cache_path = get_cache_dir() / "smpl_humanoid" + if all((cache_path / f"{name}.xml").is_file() for name in SMPL_HUMANOID_VARIANTS): + config.set_model_path("smpl-humanoid", cache_path) + return validate_path(cache_path) + return download_model() + + +def download_model() -> Path: + cache_dir = get_cache_dir() / "smpl_humanoid" + print(f"Downloading SMPL humanoid model to {cache_dir}...") + download_and_extract(url=SMPL_HUMANOID_URL, dest=cache_dir) + print("Done") + path = validate_path(cache_dir) + config.set_model_path("smpl-humanoid", path) + return path + + def _model_source(source: PathLike) -> Path: if isinstance(source, str): name = source.strip().lower().replace("-", "_") - if name in SMPL_HUMANOID_SOURCES: - return SMPL_HUMANOID_SOURCES[name] + if name in SMPL_HUMANOID_VARIANTS: + return get_model_path() / f"{name}.xml" path = Path(source) - if path.is_file(): - return path - if not path.parent.parts: + if path.parent == Path(".") and path.suffix == "": variants = ", ".join(SMPL_HUMANOID_VARIANTS) raise ValueError(f"Unknown SMPL humanoid source {source!r}. Available sources: {variants}") + else: + path = Path(source) - return Path(source) + if not path.is_file(): + raise FileNotFoundError(f"SMPL humanoid XML not found: {path}") + if path.suffix.lower() != ".xml": + raise ValueError(f"Expected an SMPL humanoid XML file, got: {path}") + return path def _walk_xml_bodies( @@ -137,7 +189,7 @@ def _walk_xml_bodies( _walk_xml_bodies(child, parent_name=parent_name, bodies=bodies, parents=parents) -def _load_xml_geoms(bodies: dict[str, ET.Element], *, dtype) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]: +def _load_xml_geoms(bodies: dict[str, ET.Element], *, dtype) -> tuple[np.ndarray, np.ndarray, LinkData]: vertices_by_link = [] faces_by_link = [] joint_indices = [] @@ -171,17 +223,17 @@ def _load_xml_geoms(bodies: dict[str, ET.Element], *, dtype) -> tuple[np.ndarray if not vertices_by_link: raise ValueError("SMPL humanoid XML does not contain any primitive geoms.") - link_data = { - "joint_indices": joint_indices, - "vertex_starts": vertex_starts, - "vertex_counts": vertex_counts, - "face_starts": face_starts, - "face_counts": face_counts, - "geom_positions": np.asarray(geom_positions, dtype=dtype), - "geom_rotations": np.asarray(geom_rotations, dtype=dtype), - "names": names, - } - return np.concatenate(vertices_by_link), np.concatenate(faces_by_link), link_data + links = LinkData( + joint_indices=joint_indices, + vertex_starts=vertex_starts, + vertex_counts=vertex_counts, + face_starts=face_starts, + face_counts=face_counts, + geom_positions=np.asarray(geom_positions, dtype=dtype), + geom_rotations=np.asarray(geom_rotations, dtype=dtype), + names=names, + ) + return np.concatenate(vertices_by_link), np.concatenate(faces_by_link), links def _geom_mesh(geom: ET.Element, *, dtype) -> tuple[np.ndarray, np.ndarray]: @@ -236,7 +288,6 @@ def _mesh_arrays(mesh: Trimesh, *, dtype) -> tuple[np.ndarray, np.ndarray]: __all__ = [ - "SMPL_HUMANOID_SOURCES", "SmplHumanoidWeights", "load_model_data", ] diff --git a/src/body_models/skeletons/myofullbody/io.py b/src/body_models/skeletons/myofullbody/io.py index 27b3e44f..c446ccaf 100644 --- a/src/body_models/skeletons/myofullbody/io.py +++ b/src/body_models/skeletons/myofullbody/io.py @@ -17,7 +17,7 @@ from jaxtyping import Float, Int from body_models import config -from body_models.cache import download_and_extract, get_cache_dir +from body_models.cache import HF_MODEL_BASE_URL, download_and_extract, get_cache_dir from body_models.common.stl import load_stl_mesh as _load_stl_mesh from body_models.robots import mjcf @@ -28,7 +28,7 @@ _RY_90 = np.array([[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]], dtype=np.float32) _MUJOCO_TO_KIMODO_BARE = np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float32) MUJOCO_TO_KIMODO = (_RY_90 @ _MUJOCO_TO_KIMODO_BARE).astype(np.float32) -MUSCLEMIMIC_REPO_ZIP = "https://github.com/amathislab/musclemimic_models/archive/refs/heads/main.zip" +MUSCLEMIMIC_REPO_ZIP = f"{HF_MODEL_BASE_URL}/myofullbody/assets.zip" MAIN_XML_RELPATH = Path("body") / "myofullbody.xml" ROOT_BODY_NAME = "Full Body" Array = Any @@ -93,7 +93,6 @@ def download_model() -> Path: download_and_extract( url=MUSCLEMIMIC_REPO_ZIP, dest=cache_dir, - extract_subdir="musclemimic_models-main/musclemimic_models/model/", ) print("Done") return cache_dir