SMPL-H shape fitting

SMPL-H shape fitting#

Optimizing SMPL-H body shape parameters to achieve a target height.

Inputs: Target height, initial shape parameters (zeros)
Outputs: Shape (beta) parameters that produce a body with the desired height

Features used:

  • Var for shape parameters

  • @jaxls.Cost.factory with constraint for height

  • Augmented Lagrangian solver for constrained optimization

Note: the SMPL-H implementation here is minimal. For full-featured SMPL models in jaxls, see egoallo or VideoMimic.

Hide code cell source

 
import sys
from loguru import logger

logger.remove()
logger.add(sys.stdout, format="<level>{level: <8}</level> | {message}");

import io
import pathlib
import urllib.request
import zipfile

import jax
import jax.numpy as jnp
import jax_dataclasses as jdc
import jaxls
import numpy as np
from jax import Array

Download SMPL-H model#

The SMPL-H model represents human body shape using a low-dimensional parameterization. Shape variations are controlled by beta parameters that deform a template mesh.

# Download SMPL-H model if not already present.
smplh_path = pathlib.Path("/tmp/SMPLH_NEUTRAL.npz")

if not smplh_path.exists():
    print("Downloading SMPL-H model...")
    url = "https://brentyi.github.io/viser-example-assets/SMPLH_NEUTRAL.zip"
    with urllib.request.urlopen(url) as response:
        zip_data = io.BytesIO(response.read())
    with zipfile.ZipFile(zip_data) as zf:
        zf.extractall("/tmp")
    print(f"Downloaded to {smplh_path}")
else:
    print(f"Using cached model at {smplh_path}")

Using cached model at /tmp/SMPLH_NEUTRAL.npz

SMPL-H model implementation#

A minimal implementation of the SMPL-H body model. Shape is controlled by beta parameters, which are PCA coefficients that linearly combine learned shape basis vectors to deform the template mesh.

@jdc.pytree_dataclass
class SmplhModel:
    """SMPL-H human body model."""

    faces: Array
    """Vertex indices for mesh faces, shape (faces, 3)."""
    v_template: Array
    """Template mesh vertices, shape (verts, 3)."""
    shapedirs: Array
    """Shape blend shape bases, shape (verts, 3, n_betas)."""

    @staticmethod
    def load(npz_path: pathlib.Path) -> "SmplhModel":
        """Load model from .npz file."""
        params = np.load(npz_path, allow_pickle=True)
        return SmplhModel(
            faces=jnp.array(params["f"].astype(np.int32)),
            v_template=jnp.array(params["v_template"].astype(np.float32)),
            shapedirs=jnp.array(params["shapedirs"].astype(np.float32)),
        )

    def get_vertices(self, betas: Array) -> Array:
        """Compute mesh vertices for given shape parameters."""
        num_betas = betas.shape[0]
        # Apply shape blend shapes: v = v_template + shapedirs @ betas.
        return self.v_template + jnp.einsum(
            "vxb,b->vx", self.shapedirs[:, :, :num_betas], betas
        )

    def get_height(self, betas: Array) -> Array:
        """Compute body height from min to max vertex z-coordinate."""
        verts = self.get_vertices(betas)
        # Height is the range of the y-coordinate (SMPL uses y-up).
        return jnp.max(verts[:, 1]) - jnp.min(verts[:, 1])

# Load the model.
model = SmplhModel.load(smplh_path)

# Check the template (zero-beta) height.
template_height = float(model.get_height(jnp.zeros(16)))
print(
    f"Template mesh: {model.v_template.shape[0]} vertices, {model.faces.shape[0]} faces"
)
print(f"Template height (beta=0): {template_height:.3f} m")

Template mesh: 6890 vertices, 13776 faces
Template height (beta=0): 1.717 m

Problem setup#

We optimize the first 10 beta parameters to achieve a target height of 2.0 meters (tall), while regularizing betas toward zero to maintain a natural body shape.

# Target height in meters.
TARGET_HEIGHT = 2.0
NUM_BETAS = 10


# Variable for shape parameters.
class BetaVar(jaxls.Var[jax.Array], default_factory=lambda: jnp.zeros(NUM_BETAS)):
    """SMPL-H beta (shape) parameters."""


beta_var = BetaVar(id=0)


@jaxls.Cost.factory(kind="constraint_eq_zero")
def height_constraint(
    vals: jaxls.VarValues,
    var: BetaVar,
    model: SmplhModel,
    target_height: float,
) -> jax.Array:
    """Constrain body height to target value."""
    betas = vals[var]
    current_height = model.get_height(betas)
    # jaxls accepts scalar residuals.
    return current_height - target_height


@jaxls.Cost.factory
def beta_regularization(
    vals: jaxls.VarValues,
    var: BetaVar,
    weight: float,
) -> jax.Array:
    """Regularize betas toward zero for natural shapes."""
    return weight * vals[var]

Solving#

When constraints are present, jaxls automatically uses an Augmented Lagrangian method. The solver iteratively adjusts Lagrange multipliers and penalty parameters to satisfy the constraint.

# Build the optimization problem.
costs: list[jaxls.Cost] = [
    height_constraint(beta_var, model, TARGET_HEIGHT),
    beta_regularization(beta_var, weight=0.5),
]

# Initial values: zeros.
initial_betas = jnp.zeros(NUM_BETAS)
initial_vals = jaxls.VarValues.make([beta_var.with_value(initial_betas)])

# Build the problem.
problem = jaxls.LeastSquaresProblem(costs, [beta_var])

# Visualize the problem structure structure.
problem.show()

# Analyze the problem and print info.
problem = problem.analyze()

print(f"Initial height: {model.get_height(initial_betas):.3f} m")
print(f"Target height: {TARGET_HEIGHT:.3f} m")

INFO     | Building optimization problem with 2 terms and 1 variables: 1 costs, 1 eq_zero, 0 leq_zero, 0 geq_zero
INFO     | Vectorizing group with 1 costs, 1 variables each: beta_regularization
INFO     | Vectorizing constraint group with 1 constraints (constraint_eq_zero), 1 variables each: augmented_height_constraint
Initial height: 1.717 m
Target height: 2.000 m

# Solve. Augmented Lagrangian is used automatically for constrained problems.
solution = problem.solve(
    initial_vals,
    linear_solver="dense_cholesky",
    termination=jaxls.TerminationConfig(cost_tolerance=1e-8),
)

optimized_betas = solution[beta_var]
final_height = model.get_height(optimized_betas)

print(f"\nOptimized height: {float(final_height):.4f} m")
print(f"Height error: {abs(float(final_height) - TARGET_HEIGHT) * 100:.2f} cm")
print(f"Beta norm: {float(jnp.linalg.norm(optimized_betas)):.3f}")

INFO     | Augmented Lagrangian: initial snorm=2.8264e-01, csupn=2.8264e-01, max_rho=1.0000e+01, constraint_dim=1
INFO     |  step #0: cost=0.0000 lambd=0.0005
INFO     |      - beta_regularization(1): 0.00000 (avg 0.00000)
INFO     |      - augmented_height_constraint(1): 0.79885 (avg 0.79885)
INFO     |      accepted=True ATb_norm=4.61e-01 cost_prev=0.7988 cost_new=0.5705
INFO     |  step #1: cost=0.1629 lambd=0.0003
INFO     |      - beta_regularization(1): 0.16292 (avg 0.01629)
INFO     |      - augmented_height_constraint(1): 0.40754 (avg 0.40754)
INFO     |  step #2: cost=0.1629 lambd=0.0005
INFO     |      - beta_regularization(1): 0.16292 (avg 0.01629)
INFO     |      - augmented_height_constraint(1): 0.40754 (avg 0.40754)
INFO     |  step #3: cost=0.1629 lambd=0.0010
INFO     |      - beta_regularization(1): 0.16292 (avg 0.01629)
INFO     |      - augmented_height_constraint(1): 0.40754 (avg 0.40754)
INFO     |  step #4: cost=0.1629 lambd=0.0020
INFO     |      - beta_regularization(1): 0.16292 (avg 0.01629)
INFO     |      - augmented_height_constraint(1): 0.40754 (avg 0.40754)
INFO     |      accepted=True ATb_norm=2.47e-04 cost_prev=0.5705 cost_new=0.5705
INFO     |  AL update: snorm=2.0183e-01, csupn=2.0183e-01, max_rho=4.0000e+01
INFO     |  step #5: cost=0.1631 lambd=0.0010
INFO     |      - beta_regularization(1): 0.16309 (avg 0.01631)
INFO     |      - augmented_height_constraint(1): 2.54604 (avg 2.54604)
INFO     |      accepted=True ATb_norm=1.32e+00 cost_prev=2.7091 cost_new=1.6977
INFO     |  step #6: cost=1.0495 lambd=0.0005
INFO     |      - beta_regularization(1): 1.04950 (avg 0.10495)
INFO     |      - augmented_height_constraint(1): 0.64819 (avg 0.64819)
INFO     |      accepted=True ATb_norm=1.91e-02 cost_prev=1.6977 cost_new=1.6972
INFO     |  step #7: cost=1.0559 lambd=0.0003
INFO     |      - beta_regularization(1): 1.05589 (avg 0.10559)
INFO     |      - augmented_height_constraint(1): 0.64128 (avg 0.64128)
INFO     |      accepted=True ATb_norm=3.82e-03 cost_prev=1.6972 cost_new=1.6972
INFO     |  AL update: snorm=7.5888e-02, csupn=7.5888e-02, max_rho=4.0000e+01
INFO     |  step #8: cost=1.0586 lambd=0.0001
INFO     |      - beta_regularization(1): 1.05863 (avg 0.10586)
INFO     |      - augmented_height_constraint(1): 1.63594 (avg 1.63594)
INFO     |      accepted=True ATb_norm=5.04e-01 cost_prev=2.6946 cost_new=2.5507
INFO     |  step #9: cost=1.5911 lambd=0.0001
INFO     |      - beta_regularization(1): 1.59110 (avg 0.15911)
INFO     |      - augmented_height_constraint(1): 0.95959 (avg 0.95959)
INFO     |      accepted=True ATb_norm=4.95e-03 cost_prev=2.5507 cost_new=2.5507
INFO     |  AL update: snorm=2.8384e-02, csupn=2.8384e-02, max_rho=4.0000e+01
INFO     |  step #10: cost=1.5930 lambd=0.0000
INFO     |      - beta_regularization(1): 1.59301 (avg 0.15930)
INFO     |      - augmented_height_constraint(1): 1.34123 (avg 1.34123)
INFO     |      accepted=True ATb_norm=1.89e-01 cost_prev=2.9342 cost_new=2.9141
INFO     |  step #11: cost=1.8200 lambd=0.0000
INFO     |      - beta_regularization(1): 1.81999 (avg 0.18200)
INFO     |      - augmented_height_constraint(1): 1.09412 (avg 1.09412)
INFO     |      accepted=True ATb_norm=4.14e-06 cost_prev=2.9141 cost_new=2.9141
INFO     |  AL update: snorm=1.0657e-02, csupn=1.0657e-02, max_rho=4.0000e+01
INFO     |  step #12: cost=1.8200 lambd=0.0000
INFO     |      - beta_regularization(1): 1.82000 (avg 0.18200)
INFO     |      - augmented_height_constraint(1): 1.23965 (avg 1.23965)
INFO     |      accepted=True ATb_norm=7.09e-02 cost_prev=3.0597 cost_new=3.0568
INFO     |  step #13: cost=1.9091 lambd=0.0000
INFO     |      - beta_regularization(1): 1.90912 (avg 0.19091)
INFO     |      - augmented_height_constraint(1): 1.14769 (avg 1.14769)
INFO     |      accepted=True ATb_norm=4.97e-07 cost_prev=3.0568 cost_new=3.0568
INFO     |  AL update: snorm=4.0011e-03, csupn=4.0011e-03, max_rho=4.0000e+01
INFO     |  step #14: cost=1.9091 lambd=0.0000
INFO     |      - beta_regularization(1): 1.90913 (avg 0.19091)
INFO     |      - augmented_height_constraint(1): 1.20255 (avg 1.20255)
INFO     |      accepted=True ATb_norm=2.66e-02 cost_prev=3.1117 cost_new=3.1113
INFO     |  step #15: cost=1.9431 lambd=0.0000
INFO     |      - beta_regularization(1): 1.94314 (avg 0.19431)
INFO     |      - augmented_height_constraint(1): 1.16814 (avg 1.16814)
INFO     |      accepted=False ATb_norm=1.47e-08 cost_prev=3.1113 cost_new=3.1113
INFO     |  AL update: snorm=1.5023e-03, csupn=1.5023e-03, max_rho=4.0000e+01
INFO     |  step #16: cost=1.9431 lambd=0.0000
INFO     |      - beta_regularization(1): 1.94314 (avg 0.19431)
INFO     |      - augmented_height_constraint(1): 1.18877 (avg 1.18877)
INFO     |      accepted=True ATb_norm=9.99e-03 cost_prev=3.1319 cost_new=3.1318
INFO     |  AL update: snorm=5.6410e-04, csupn=5.6410e-04, max_rho=4.0000e+01
INFO     |  step #17: cost=1.9560 lambd=0.0000
INFO     |      - beta_regularization(1): 1.95599 (avg 0.19560)
INFO     |      - augmented_height_constraint(1): 1.18361 (avg 1.18361)
INFO     |      accepted=True ATb_norm=3.75e-03 cost_prev=3.1396 cost_new=3.1396
INFO     |  AL update: snorm=2.1172e-04, csupn=2.1172e-04, max_rho=4.0000e+01
INFO     |  step #18: cost=1.9608 lambd=0.0000
INFO     |      - beta_regularization(1): 1.96082 (avg 0.19608)
INFO     |      - augmented_height_constraint(1): 1.18168 (avg 1.18168)
INFO     |      accepted=True ATb_norm=1.41e-03 cost_prev=3.1425 cost_new=3.1425
INFO     |  AL update: snorm=7.9513e-05, csupn=7.9513e-05, max_rho=4.0000e+01
INFO     |  step #19: cost=1.9626 lambd=0.0000
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #20: cost=1.9626 lambd=0.0000
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #21: cost=1.9626 lambd=0.0000
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #22: cost=1.9626 lambd=0.0001
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #23: cost=1.9626 lambd=0.0002
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #24: cost=1.9626 lambd=0.0003
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #25: cost=1.9626 lambd=0.0006
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #26: cost=1.9626 lambd=0.0013
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #27: cost=1.9626 lambd=0.0026
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #28: cost=1.9626 lambd=0.0051
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #29: cost=1.9626 lambd=0.0102
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #30: cost=1.9626 lambd=0.0205
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #31: cost=1.9626 lambd=0.0410
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #32: cost=1.9626 lambd=0.0819
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |  step #33: cost=1.9626 lambd=0.1638
INFO     |      - beta_regularization(1): 1.96264 (avg 0.19626)
INFO     |      - augmented_height_constraint(1): 1.18095 (avg 1.18095)
INFO     |      accepted=True ATb_norm=5.29e-04 cost_prev=3.1436 cost_new=3.1436
INFO     |  AL update: snorm=3.4094e-05, csupn=3.4094e-05, max_rho=4.0000e+01
INFO     |  step #34: cost=1.9633 lambd=0.0819
INFO     |      - beta_regularization(1): 1.96326 (avg 0.19633)
INFO     |      - augmented_height_constraint(1): 1.18080 (avg 1.18080)
INFO     |  step #35: cost=1.9633 lambd=0.1638
INFO     |      - beta_regularization(1): 1.96326 (avg 0.19633)
INFO     |      - augmented_height_constraint(1): 1.18080 (avg 1.18080)
INFO     |      accepted=True ATb_norm=2.72e-04 cost_prev=3.1441 cost_new=3.1441
INFO     |  AL update: snorm=1.0848e-05, csupn=1.0848e-05, max_rho=4.0000e+01
INFO     |  step #36: cost=1.9636 lambd=0.0819
INFO     |      - beta_regularization(1): 1.96358 (avg 0.19636)
INFO     |      - augmented_height_constraint(1): 1.18063 (avg 1.18063)
INFO     |      accepted=True ATb_norm=9.59e-05 cost_prev=3.1442 cost_new=3.1442
INFO     |  AL update: snorm=2.1458e-06, csupn=2.1458e-06, max_rho=4.0000e+01
INFO     | Terminated @ iteration #37: cost=1.9637 criteria=[0 1 0], term_deltas=6.0e-05,9.4e-05,3.0e-05

Optimized height: 2.0000 m
Height error: 0.00 cm
Beta norm: 2.803

Visualization#

Compare the template mesh (beta=0) with the optimized shape side by side.

Hide code cell source

 
import contextlib
import io
import viser

# Get vertices for both configurations.
initial_verts = np.array(model.get_vertices(initial_betas))
optimized_verts = np.array(model.get_vertices(optimized_betas))
faces = np.array(model.faces)

# Compute heights for labels.
initial_height = float(model.get_height(initial_betas))
optimized_height = float(final_height)

# Compute y offsets to align feet at ground level (SMPL uses y-up).
initial_y_offset = -initial_verts[:, 1].min()
optimized_y_offset = -optimized_verts[:, 1].min()

# Offset for side-by-side placement.
x_offset = 0.8

# Create Viser server (suppress output).
try:
    server.stop()
except NameError:
    pass
with (
    contextlib.redirect_stdout(io.StringIO()),
    contextlib.redirect_stderr(io.StringIO()),
):
    server = viser.ViserServer(verbose=False)

# Use y-up to match SMPL coordinates directly.
server.scene.set_up_direction("+y")

# Set initial camera position in front of the figures (SMPL faces +z).
server.initial_camera.position = (0.0, 2.0, 4.0)
server.initial_camera.look_at = (0.0, 0.9, 0.0)

# Add ground grid (xz plane for y-up) with fade.
server.scene.add_grid(
    "/ground",
    width=4.0,
    height=2.0,
    plane="xz",
    infinite_grid=True,
    fade_distance=10.0,
    cell_color=(200, 200, 200),
    section_color=(170, 170, 170),
)

# Add initial mesh (template, on the right).
server.scene.add_mesh_simple(
    "/initial_mesh",
    vertices=initial_verts + np.array([[-x_offset, initial_y_offset, 0.0]]),
    faces=faces,
    color=(70, 130, 180),  # Steel blue
    flat_shading=False,
)

# Add optimized mesh (tall figure, on the left).
server.scene.add_mesh_simple(
    "/optimized_mesh",
    vertices=optimized_verts + np.array([[x_offset, optimized_y_offset, 0.0]]),
    faces=faces,
    color=(34, 139, 34),  # Forest green
    flat_shading=False,
)

# Add height labels above each mesh.
server.scene.add_label(
    "/initial_label",
    text=f"Template: {initial_height:.2f}m",
    position=(-x_offset, initial_height + 0.15, 0.0),
    anchor="bottom-center",
)
server.scene.add_label(
    "/optimized_label",
    text=f"Optimized: {optimized_height:.2f}m",
    position=(x_offset, optimized_height + 0.15, 0.0),
    anchor="bottom-center",
)

# Display inline in the notebook.
server.scene.show(height=500)

╭────── viser (listening *:8082) ───────╮
│             ╷                         │
│   HTTP      │ http://localhost:8082   │
│   Websocket │ ws://localhost:8082     │
│             ╵                         │
╰───────────────────────────────────────╯

The optimization finds shape parameters that satisfy the height constraint while keeping the body shape natural (small beta norm). The regularization prevents extreme deformations that could produce unrealistic body shapes.

For more on constrained optimization, see Constraints.

Contents