keypoints2body 0.1.0__tar.gz

keypoints2body-0.1.0/PKG-INFO

@@ -0,0 +1,75 @@
+Metadata-Version: 2.4
+Name: keypoints2body
+Version: 0.1.0
+Summary: SMPL/SMPLH/SMPLX optimization library for single frames and sequences
+Author: keypoints2body contributors
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: torch
+Requires-Dist: smplx
+Requires-Dist: h5py
+Requires-Dist: tqdm
+Provides-Extra: dev
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: sphinx>=7; extra == "docs"
+Requires-Dist: furo>=2024.8.6; extra == "docs"
+Requires-Dist: sphinxcontrib-mermaid>=0.9; extra == "docs"
+Provides-Extra: eval
+Provides-Extra: render
+
+![keypoints2body header](docs/assets/header.svg)
+
+# keypoints2body
+
+`keypoints2body` is a Python library for optimizing SMPL-family body model parameters
+from 3D joints for both single frames and motion sequences.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+## Quick usage
+
+```python
+import numpy as np
+from keypoints2body import optimize_params_frame
+
+joints = np.zeros((22, 3), dtype=np.float32)
+result = optimize_params_frame(joints, joint_layout="AMASS")
+```
+
+## Documentation
+
+Full documentation lives under [`docs/`](docs/) and is intended for Sphinx.
+
+Suggested starting points:
+
+- [Getting started](docs/getting_started.rst)
+- [Library usage](docs/usage.rst)
+- [CLI reference](docs/cli.rst)
+- [Architecture](docs/architecture.rst)
+- [API reference](docs/api.rst)
+- [Contributor guide](docs/contributing.rst)
+
+## CLI
+
+```bash
+keypoints2body-fit-frame --help
+keypoints2body-fit-seq --help
+keypoints2body-eval --help
+```
+
+Project script:
+
+```bash
+python smpl_fit.py --help
+```

keypoints2body-0.1.0/README.md

@@ -0,0 +1,49 @@
+![keypoints2body header](docs/assets/header.svg)
+
+# keypoints2body
+
+`keypoints2body` is a Python library for optimizing SMPL-family body model parameters
+from 3D joints for both single frames and motion sequences.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+## Quick usage
+
+```python
+import numpy as np
+from keypoints2body import optimize_params_frame
+
+joints = np.zeros((22, 3), dtype=np.float32)
+result = optimize_params_frame(joints, joint_layout="AMASS")
+```
+
+## Documentation
+
+Full documentation lives under [`docs/`](docs/) and is intended for Sphinx.
+
+Suggested starting points:
+
+- [Getting started](docs/getting_started.rst)
+- [Library usage](docs/usage.rst)
+- [CLI reference](docs/cli.rst)
+- [Architecture](docs/architecture.rst)
+- [API reference](docs/api.rst)
+- [Contributor guide](docs/contributing.rst)
+
+## CLI
+
+```bash
+keypoints2body-fit-frame --help
+keypoints2body-fit-seq --help
+keypoints2body-eval --help
+```
+
+Project script:
+
+```bash
+python smpl_fit.py --help
+```

keypoints2body-0.1.0/keypoints2body/__init__.py

@@ -0,0 +1,28 @@
+from importlib.metadata import PackageNotFoundError, version
+
+from .api.frame import optimize_params_frame
+from .api.sequence import optimize_params_sequence, optimize_shape_sequence
+from .models.smpl_data import (
+    BodyModelFitResult,
+    BodyModelParams,
+    SMPLData,
+    SMPLHData,
+    SMPLXData,
+)
+
+try:
+    __version__ = version("keypoints2body")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+__all__ = [
+    "__version__",
+    "optimize_params_frame",
+    "optimize_params_sequence",
+    "optimize_shape_sequence",
+    "BodyModelFitResult",
+    "BodyModelParams",
+    "SMPLData",
+    "SMPLHData",
+    "SMPLXData",
+]

keypoints2body-0.1.0/keypoints2body/api/__init__.py

@@ -0,0 +1,8 @@
+from .frame import optimize_params_frame
+from .sequence import optimize_params_sequence, optimize_shape_sequence
+
+__all__ = [
+    "optimize_params_frame",
+    "optimize_params_sequence",
+    "optimize_shape_sequence",
+]

keypoints2body-0.1.0/keypoints2body/api/frame.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ..core.config import BodyModelConfig, FrameOptimizeConfig
+from ..core.engine import OptimizeEngine, default_init_params, load_mean_pose_shape
+from ..core.joints.adapters import adapt_layout_and_conf, normalize_joints_frame
+from ..models.smpl_data import BodyModelFitResult, BodyModelParams, SMPLData
+from .model_factory import load_body_model
+
+DEFAULT_MEAN_FILE = "./data/models/neutral_smpl_mean_params.h5"
+
+
+def optimize_params_frame(
+    joints,
+    *,
+    prev_params: Optional[BodyModelParams] = None,
+    body_model: str = "smpl",
+    joint_layout: Optional[str] = None,
+    model=None,
+    config: Optional[FrameOptimizeConfig | dict] = None,
+    device=None,
+) -> BodyModelFitResult:
+    """Optimize body parameters for a single frame of 3D joints.
+
+    Args:
+        joints: Frame keypoints with shape ``(K,3)`` or ``(K,4)``.
+        prev_params: Optional warm-start parameters from a previous fit.
+        body_model: Body model type (for example ``"smpl"``).
+        joint_layout: Optional explicit layout label for adapter selection.
+        model: Optional preloaded body model instance.
+        config: Optional frame config or dict equivalent.
+        device: Optional torch device specifier.
+
+    Returns:
+        Typed fitting result containing optimized parameters, joints, vertices, and loss.
+    """
+    device = torch.device(device) if device is not None else torch.device("cpu")
+
+    if isinstance(config, dict):
+        frame_cfg = FrameOptimizeConfig(**config)
+    elif isinstance(config, FrameOptimizeConfig):
+        frame_cfg = config
+    else:
+        frame_cfg = FrameOptimizeConfig()
+
+    if frame_cfg.input_type != "joints3d":
+        raise NotImplementedError(
+            f"input_type='{frame_cfg.input_type}' is not implemented in this release. "
+            "Current APIs support only joints3d."
+        )
+
+    j3d, conf_3d = normalize_joints_frame(joints)
+    j_np = j3d.squeeze(0).cpu().numpy()[None, ...]
+    c_np = conf_3d.cpu().numpy()[None, ...]
+    j_np, c_np, out_layout = adapt_layout_and_conf(j_np, c_np, joint_layout)
+    j3d = torch.as_tensor(j_np, dtype=torch.float32, device=device)
+    conf_3d = torch.as_tensor(c_np[0], dtype=torch.float32, device=device)
+
+    if out_layout not in ("SMPL24", "AMASS"):
+        raise ValueError(f"Unsupported output layout after adaptation: {out_layout}")
+    frame_cfg.joints_category = out_layout
+
+    if model is None:
+        body_cfg = BodyModelConfig(model_type=body_model)
+        model = load_body_model(body_cfg, device)
+
+    engine = OptimizeEngine(model=model, frame_config=frame_cfg, device=device)
+
+    if prev_params is None:
+        init_mean_pose, init_mean_shape = load_mean_pose_shape(
+            DEFAULT_MEAN_FILE, device
+        )
+        init_params = default_init_params(
+            init_mean_pose,
+            init_mean_shape,
+            j3d,
+            model,
+            joints_category=frame_cfg.joints_category,
+            coordinate_mode=frame_cfg.coordinate_mode,
+        )
+    else:
+        init_params = prev_params.to(device)
+        if frame_cfg.coordinate_mode == "world" and init_params.transl is None:
+            pose = init_params.pose
+            if not isinstance(pose, torch.Tensor):
+                pose = torch.as_tensor(pose, dtype=torch.float32, device=device)
+            betas = init_params.betas
+            if not isinstance(betas, torch.Tensor):
+                betas = torch.as_tensor(betas, dtype=torch.float32, device=device)
+            transl = default_init_params(
+                pose,
+                betas,
+                j3d,
+                model,
+                frame_cfg.joints_category,
+                frame_cfg.coordinate_mode,
+            ).transl
+            init_params = SMPLData(
+                betas=betas,
+                global_orient=pose[:, :3],
+                body_pose=pose[:, 3:],
+                transl=transl,
+                metadata=dict(getattr(init_params, "metadata", {})),
+            )
+
+    return engine.fit_frame(
+        init_params=init_params, j3d=j3d, conf_3d=conf_3d, seq_ind=0
+    )

keypoints2body-0.1.0/keypoints2body/api/model_factory.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import smplx
+import torch
+
+from ..core.config import BodyModelConfig
+
+
+def load_body_model(
+    config: BodyModelConfig,
+    device: torch.device,
+):
+    """Load a body model module with ``smplx.create``.
+
+    Args:
+        config: Model-loading configuration.
+        device: Torch device where the model should live.
+
+    Returns:
+        A body model module returned by ``smplx.create`` moved to ``device``.
+    """
+    model_dir = Path(config.model_dir).expanduser()
+    return smplx.create(
+        str(model_dir),
+        model_type=config.model_type,
+        gender=config.gender,
+        ext=config.ext,
+        batch_size=config.batch_size,
+    ).to(device)

keypoints2body-0.1.0/keypoints2body/api/sequence.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ..core.config import BodyModelConfig, FrameOptimizeConfig, SequenceOptimizeConfig
+from ..core.engine import (
+    OptimizeEngine,
+    default_init_params,
+    load_mean_pose_shape,
+    optimize_shape_pass,
+)
+from ..core.joints.adapters import adapt_layout_and_conf, normalize_joints_sequence
+from ..models.smpl_data import BodyModelFitResult, BodyModelParams, SMPLData
+from .model_factory import load_body_model
+
+DEFAULT_MEAN_FILE = "./data/models/neutral_smpl_mean_params.h5"
+
+
+def optimize_params_sequence(
+    joints_seq,
+    *,
+    init_params: Optional[BodyModelParams] = None,
+    body_model: str = "smpl",
+    joint_layout: Optional[str] = None,
+    model=None,
+    config: Optional[SequenceOptimizeConfig | dict] = None,
+    device=None,
+) -> list[BodyModelFitResult]:
+    """Optimize body parameters for a full motion sequence.
+
+    Args:
+        joints_seq: Sequence keypoints with shape ``(T,K,3)`` or ``(T,K,4)``.
+        init_params: Optional initial parameters for frame 0.
+        body_model: Body model type (for example ``"smpl"``).
+        joint_layout: Optional explicit layout label for adapter selection.
+        model: Optional preloaded body model instance.
+        config: Optional sequence config or dict equivalent.
+        device: Optional torch device specifier.
+
+    Returns:
+        Per-frame optimization results in temporal order.
+    """
+    device = torch.device(device) if device is not None else torch.device("cpu")
+
+    if isinstance(config, dict):
+        frame_cfg = FrameOptimizeConfig(**config.get("frame", {}))
+        seq_cfg = SequenceOptimizeConfig(
+            frame=frame_cfg,
+            num_shape_iters=config.get("num_shape_iters", 40),
+            num_shape_frames=config.get("num_shape_frames", 50),
+            use_shape_optimization=config.get("use_shape_optimization", True),
+            use_previous_frame_init=config.get("use_previous_frame_init", True),
+            fix_foot=config.get("fix_foot", False),
+            limit_frames=config.get("limit_frames", None),
+        )
+    elif isinstance(config, SequenceOptimizeConfig):
+        seq_cfg = config
+    else:
+        seq_cfg = SequenceOptimizeConfig()
+
+    if seq_cfg.frame.input_type != "joints3d":
+        raise NotImplementedError(
+            f"input_type='{seq_cfg.frame.input_type}' is not implemented in this release. "
+            "Current APIs support only joints3d."
+        )
+
+    xyz, conf = normalize_joints_sequence(joints_seq)
+    xyz_np, conf_np, out_layout = adapt_layout_and_conf(
+        xyz.cpu().numpy(), conf.cpu().numpy(), joint_layout
+    )
+    xyz = torch.as_tensor(xyz_np, dtype=torch.float32, device=device)
+    conf = torch.as_tensor(conf_np, dtype=torch.float32, device=device)
+
+    if out_layout not in ("SMPL24", "AMASS"):
+        raise ValueError(f"Unsupported output layout after adaptation: {out_layout}")
+    seq_cfg.frame.joints_category = out_layout
+
+    if seq_cfg.limit_frames is not None and seq_cfg.limit_frames > 0:
+        xyz = xyz[: seq_cfg.limit_frames]
+        conf = conf[: seq_cfg.limit_frames]
+
+    if seq_cfg.fix_foot and xyz.shape[1] > 11:
+        conf[:, 7] = 1.5
+        conf[:, 8] = 1.5
+        conf[:, 10] = 1.5
+        conf[:, 11] = 1.5
+
+    if model is None:
+        body_cfg = BodyModelConfig(model_type=body_model)
+        model = load_body_model(body_cfg, device)
+
+    init_mean_pose, init_mean_shape = load_mean_pose_shape(DEFAULT_MEAN_FILE, device)
+    betas_opt = optimize_shape_pass(
+        model=model,
+        seq_config=seq_cfg,
+        init_mean_shape=init_mean_shape,
+        init_mean_pose=init_mean_pose,
+        data_tensor=xyz,
+        confidence_input=conf[0],
+        device=device,
+    )
+
+    engine = OptimizeEngine(model=model, frame_config=seq_cfg.frame, device=device)
+    results: list[BodyModelFitResult] = []
+
+    if init_params is None:
+        prev = default_init_params(
+            init_mean_pose,
+            betas_opt,
+            xyz[0:1],
+            model,
+            joints_category=seq_cfg.frame.joints_category,
+            coordinate_mode=seq_cfg.frame.coordinate_mode,
+        )
+    else:
+        prev = init_params.to(device)
+
+    for idx in range(xyz.shape[0]):
+        frame = xyz[idx : idx + 1]
+        frame_conf = conf[idx]
+
+        if seq_cfg.frame.coordinate_mode == "world" and prev.transl is None:
+            pose = (
+                prev.pose
+                if isinstance(prev.pose, torch.Tensor)
+                else torch.as_tensor(prev.pose, dtype=torch.float32, device=device)
+            )
+            betas = (
+                prev.betas
+                if isinstance(prev.betas, torch.Tensor)
+                else torch.as_tensor(prev.betas, dtype=torch.float32, device=device)
+            )
+            prev = SMPLData(
+                betas=betas,
+                global_orient=pose[:, :3],
+                body_pose=pose[:, 3:],
+                transl=default_init_params(
+                    pose,
+                    betas,
+                    frame,
+                    model,
+                    seq_cfg.frame.joints_category,
+                    seq_cfg.frame.coordinate_mode,
+                ).transl,
+                metadata=dict(getattr(prev, "metadata", {})),
+            )
+
+        res = engine.fit_frame(
+            init_params=prev, j3d=frame, conf_3d=frame_conf, seq_ind=idx
+        )
+        results.append(res)
+        if seq_cfg.use_previous_frame_init:
+            prev = res.params
+
+    return results
+
+
+def optimize_shape_sequence(
+    joints_seq,
+    *,
+    body_model: str = "smpl",
+    joint_layout: Optional[str] = None,
+    model=None,
+    config: Optional[SequenceOptimizeConfig | dict] = None,
+    device=None,
+) -> BodyModelParams:
+    """Run sequence optimization and return the final frame parameters.
+
+    Args:
+        joints_seq: Sequence keypoints input.
+        body_model: Body model type.
+        joint_layout: Optional explicit layout label.
+        model: Optional preloaded body model instance.
+        config: Optional sequence config.
+        device: Optional torch device specifier.
+
+    Returns:
+        Parameter object from the last optimized frame.
+    """
+    results = optimize_params_sequence(
+        joints_seq,
+        init_params=None,
+        body_model=body_model,
+        joint_layout=joint_layout,
+        model=model,
+        config=config,
+        device=device,
+    )
+    if not results:
+        raise ValueError("No frames were optimized")
+    return results[-1].params

keypoints2body-0.1.0/keypoints2body/cli/__init__.py

@@ -0,0 +1 @@
+__all__ = []