Add emotion_to_mesh and bs_to_mesh PLS visualization models

feat/plotting.py

@@ -18,7 +18,7 @@
     mask_image,
     procrustes_align_2d_batched,
 )
-from feat.utils import flatten_list
+from feat.utils import flatten_list, FEAT_EMOTION_COLUMNS, MP_BLENDSHAPE_NAMES
 from huggingface_hub import hf_hub_download
 from math import sin, cos
 import warnings
@@ -55,6 +55,9 @@
     "emotion_annotation_position",
     "load_face_mesh_viz_model",
     "predict_face_mesh",
+    "load_emotion_face_mesh_model",
+    "load_blendshape_face_mesh_model",
+    "predict_face_mesh_from_features",
     "plot_face_mesh",
     "plot_face_mesh_plotly",
 ]
@@ -1873,6 +1876,13 @@ def load_face_mesh_viz_model(verbose=False, model_version="v5"):
                                             model_version=model_version)
 
 
+def _flat_to_mesh(flat, is_single):
+    """Reshape (n, 1434) axis-major [x|y|z] flat coords to (n, 478, 3), or
+    (478, 3) when ``is_single``."""
+    mesh = np.stack([flat[:, :478], flat[:, 478:956], flat[:, 956:]], axis=-1)
+    return mesh[0] if is_single else mesh
+
+
 def predict_face_mesh(au, model=None):
     """Predict the 3D MediaPipe FaceMesh from AU intensities.
 
@@ -1902,11 +1912,146 @@ def predict_face_mesh(au, model=None):
             f"au vector must be length {model.n_components}; got {au_arr.shape[1]}."
         )
     flat = model.predict(au_arr)  # (n, 1434), axis-major [x | y | z]
-    xs = flat[:, :478]
-    ys = flat[:, 478:956]
-    zs = flat[:, 956:]
-    mesh = np.stack([xs, ys, zs], axis=-1)  # (n, 478, 3)
-    return mesh[0] if is_single else mesh
+    return _flat_to_mesh(flat, is_single)
+
+
+# ---------------------------------------------------------------------
+# Emotion / blendshape → 478-vertex MediaPipe FaceMesh.
+# Companions to the AU→mesh model above, trained on the same CelebV-HQ
+# Detectorv2/MPDetector predictions and aligned into the SAME pose-canonical
+# frame (au_to_mesh v5 anchors), so all three render coherently. Weights live
+# in the py-feat/emotion_to_mesh and py-feat/bs_to_mesh HF Hub repos.
+# ---------------------------------------------------------------------
+
+_FEAT_MESH_SPECS = {
+    "emotion": ("py-feat/emotion_to_mesh", "emotion_to_mesh_pls", FEAT_EMOTION_COLUMNS),
+    "blendshape": ("py-feat/bs_to_mesh", "bs_to_mesh_pls", MP_BLENDSHAPE_NAMES),
+}
+_PLS_FEAT_MESH_MODELS = {}      # (feature, version) -> PLSFeatMeshModel
+
+
+class PLSFeatMeshModel:
+    """Wrapper around a {feature} + pose → 478-vertex MP mesh PLS (full rank).
+
+    Like ``PLSAUMeshModel`` but for a generic input feature family — 7 emotion
+    probabilities or 52 MediaPipe blendshapes. Pose is held implicit-zero at
+    inference (the absorbed feature×pose interaction terms drop out), so
+    ``predict(x)`` is a single matmul to (n, 1434) axis-major mesh coords in the
+    shared pose-canonical frame.
+    """
+
+    def __init__(self, coef, intercept, feature_columns, pose_columns,
+                 mean_aligned_mesh, feature_name, model_name):
+        self._coef = np.asarray(coef, dtype=np.float32)
+        self._intercept = np.asarray(intercept, dtype=np.float32)
+        self.feature_columns = list(feature_columns)
+        self.pose_columns = list(pose_columns)
+        self.feature_name = feature_name
+        self.n_components = len(self.feature_columns)
+        self.mean_aligned_mesh = np.asarray(mean_aligned_mesh, dtype=np.float32)
+        self.model_name_ = model_name
+
+    def predict(self, feats):
+        x_in = np.asarray(feats, dtype=np.float32)
+        if x_in.ndim == 1:
+            x_in = x_in.reshape(1, -1)
+        if x_in.ndim != 2 or x_in.shape[1] != self.n_components:
+            raise ValueError(
+                f"{self.feature_name} input must be a length-{self.n_components} "
+                f"vector or (n, {self.n_components}) batch (matching "
+                f"{self.feature_columns}); got shape {x_in.shape}."
+            )
+        # pose channels are implicit-zero, so only the leading feature rows of
+        # the deployed coef contribute — slice instead of zero-padding.
+        return x_in @ self._coef[: self.n_components] + self._intercept
+
+    def __repr__(self):
+        return (
+            f"PLSFeatMeshModel(model_name='{self.model_name_}', "
+            f"feature='{self.feature_name}', n_components={self.n_components}, "
+            f"output_shape=(n_samples, 478, 3))"
+        )
+
+
+def _load_pls_feat_to_mesh_from_hub(feature, verbose=False, model_version="v5"):
+    if feature not in _FEAT_MESH_SPECS:
+        raise ValueError(
+            f"feature must be one of {list(_FEAT_MESH_SPECS)}; got {feature!r}."
+        )
+    key = (feature, model_version)
+    if key in _PLS_FEAT_MESH_MODELS:
+        return _PLS_FEAT_MESH_MODELS[key]
+    repo_id, stem, expected = _FEAT_MESH_SPECS[feature]
+    fname = f"{stem}_{model_version}.npz"
+    if verbose:
+        print(f"Loading {feature}→mesh PLS ({model_version}) from HuggingFace Hub")
+    path = hf_hub_download(
+        repo_id=repo_id, filename=fname, cache_dir=get_resource_path(),
+    )
+    z = np.load(path, allow_pickle=False)
+    feature_columns = [str(s) for s in z["feature_columns"]]
+    if feature_columns != list(expected):
+        raise RuntimeError(
+            f"{feature}→mesh PLS feature_columns drifted. "
+            f"NPZ: {feature_columns}; expected: {list(expected)}."
+        )
+    model = PLSFeatMeshModel(
+        coef=z["coef"], intercept=z["intercept"], feature_columns=feature_columns,
+        pose_columns=[str(s) for s in z["pose_columns"]],
+        mean_aligned_mesh=z["mean_aligned_mesh"], feature_name=feature,
+        model_name=f"{stem}_{model_version}",
+    )
+    _PLS_FEAT_MESH_MODELS[key] = model
+    return model
+
+
+def load_emotion_face_mesh_model(verbose=False, model_version="v5"):
+    """Load the emotion + pose → 478-pt MediaPipe FaceMesh PLS model.
+
+    ``.predict(emotion)`` takes a length-7 vector (or ``(n, 7)`` batch) in
+    ``FEAT_EMOTION_COLUMNS`` order and returns the 478-vertex mesh (flattened
+    1434-d, axis-major) in the same pose-canonical frame as the AU→mesh model.
+    Underlying weights live in the ``py-feat/emotion_to_mesh`` HF Hub repo.
+    """
+    return _load_pls_feat_to_mesh_from_hub("emotion", verbose, model_version)
+
+
+def load_blendshape_face_mesh_model(verbose=False, model_version="v5"):
+    """Load the blendshape + pose → 478-pt MediaPipe FaceMesh PLS model.
+
+    ``.predict(blendshapes)`` takes a length-52 vector (or ``(n, 52)`` batch) in
+    ``MP_BLENDSHAPE_NAMES`` order (MPDetector output order) and returns the
+    478-vertex mesh (flattened 1434-d, axis-major) in the same pose-canonical
+    frame as the AU→mesh model. Weights live in the ``py-feat/bs_to_mesh`` HF
+    Hub repo.
+    """
+    return _load_pls_feat_to_mesh_from_hub("blendshape", verbose, model_version)
+
+
+def predict_face_mesh_from_features(feats, model):
+    """Predict the 3D MediaPipe FaceMesh from an emotion or blendshape vector.
+
+    Args:
+        feats: feature vector or batch — ``(7,)``/``(n, 7)`` for emotion,
+            ``(52,)``/``(n, 52)`` for blendshapes, matching ``model``.
+        model: a ``PLSFeatMeshModel`` from ``load_emotion_face_mesh_model()`` or
+            ``load_blendshape_face_mesh_model()``.
+
+    Returns:
+        ``(478, 3)`` for a 1-D input, ``(n, 478, 3)`` for a batch, in the
+        pose-canonical frame shared with the AU→mesh model.
+    """
+    if not isinstance(model, PLSFeatMeshModel):
+        raise ValueError(
+            "model must be a PLSFeatMeshModel (from load_emotion_face_mesh_model() "
+            "or load_blendshape_face_mesh_model())"
+        )
+    arr = np.asarray(feats)
+    is_single = arr.ndim <= 1
+    if is_single:
+        arr = arr.reshape(1, -1)
+    flat = model.predict(arr)
+    return _flat_to_mesh(flat, is_single)
 
 
 def plot_face_mesh(
@@ -1918,6 +2063,8 @@ def plot_face_mesh(
     alpha=0.9,
     view_init=(0, -90),
     *,
+    emotion=None,
+    blendshapes=None,
     mesh=None,
     mode="contours",
     gaze=None,
@@ -1967,8 +2114,11 @@ def plot_face_mesh(
     from mpl_toolkits.mplot3d import Axes3D  # noqa: F401  (registers 3d projection)
     from feat.utils.mp_plotting import FaceLandmarksConnections
 
-    if mesh is not None and au is not None:
-        raise ValueError("pass either `au` or `mesh`, not both")
+    n_given = sum(x is not None for x in (au, emotion, blendshapes, mesh))
+    if n_given > 1:
+        raise ValueError(
+            "pass at most one of `au`, `emotion`, `blendshapes`, or `mesh`"
+        )
 
     if mesh is not None:
         verts = np.asarray(mesh, dtype=np.float32)
@@ -1977,6 +2127,21 @@ def plot_face_mesh(
                 f"mesh must have shape (478, 3); got {verts.shape}. "
                 "For batched predictions, plot one face at a time."
             )
+    elif emotion is not None or blendshapes is not None:
+        feature, feats = (
+            ("emotion", emotion) if emotion is not None else ("blendshape", blendshapes)
+        )
+        if model is None:
+            model = _load_pls_feat_to_mesh_from_hub(feature)
+        feats = np.asarray(feats)
+        if feats.ndim == 2 and feats.shape[0] == 1:
+            feats = feats[0]                       # accept a single-face (1, n) row
+        verts = predict_face_mesh_from_features(feats, model=model)
+        if verts.ndim != 2:
+            raise ValueError(
+                f"plot_face_mesh expects a single {feature} vector; pass one face "
+                "at a time. For batches use predict_face_mesh_from_features()."
+            )
     elif au is None:
         if model is None:
             model = load_face_mesh_viz_model()

feat/tests/test_feat_mesh_viz.py

@@ -0,0 +1,162 @@
+"""Tests for the emotion / blendshape → 478-pt MediaPipe FaceMesh PLS models.
+
+Covers ``PLSFeatMeshModel``, ``predict_face_mesh_from_features``, the
+``load_emotion_face_mesh_model`` / ``load_blendshape_face_mesh_model`` loaders,
+and the ``emotion=`` / ``blendshapes=`` paths of ``plot_face_mesh``. Shape
+contracts run offline by stubbing the module-level cache; two
+``@pytest.mark.network`` tests load the real npz from HF Hub.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import matplotlib
+matplotlib.use("Agg")  # headless
+
+from feat import plotting as plt_mod
+from feat.plotting import (
+    PLSFeatMeshModel,
+    load_emotion_face_mesh_model,
+    load_blendshape_face_mesh_model,
+    predict_face_mesh_from_features,
+    plot_face_mesh,
+)
+from feat.utils import FEAT_EMOTION_COLUMNS, MP_BLENDSHAPE_NAMES
+
+
+def _make_stub(feature, cols):
+    rng = np.random.default_rng(0)
+    nfeat = len(cols)
+    coef = rng.standard_normal((nfeat + 3, 1434)).astype(np.float32) * 0.05
+    intercept = np.empty(1434, dtype=np.float32)
+    intercept[:478] = rng.uniform(20, 195, 478)       # x  (v5 pixel frame)
+    intercept[478:956] = rng.uniform(-247, -40, 478)  # y
+    intercept[956:] = rng.uniform(-72, 51, 478)       # z
+    mean_mesh = np.column_stack([intercept[:478], intercept[478:956], intercept[956:]])
+    return PLSFeatMeshModel(
+        coef=coef, intercept=intercept, feature_columns=cols,
+        pose_columns=["Pitch", "Yaw", "Roll"], mean_aligned_mesh=mean_mesh,
+        feature_name=feature, model_name=f"{feature}_to_mesh_pls_stub",
+    )
+
+
+@pytest.fixture
+def stub_emotion_model(monkeypatch):
+    m = _make_stub("emotion", FEAT_EMOTION_COLUMNS)
+    monkeypatch.setattr(plt_mod, "_PLS_FEAT_MESH_MODELS", {("emotion", "v5"): m})
+    return m
+
+
+@pytest.fixture
+def stub_blendshape_model(monkeypatch):
+    m = _make_stub("blendshape", MP_BLENDSHAPE_NAMES)
+    monkeypatch.setattr(plt_mod, "_PLS_FEAT_MESH_MODELS", {("blendshape", "v5"): m})
+    return m
+
+
+class TestPLSFeatMeshModel:
+    def test_emotion_predict_shape(self, stub_emotion_model):
+        flat = stub_emotion_model.predict(np.zeros((4, 7)))
+        assert flat.shape == (4, 1434)
+
+    def test_blendshape_predict_shape(self, stub_blendshape_model):
+        flat = stub_blendshape_model.predict(np.zeros((3, 52)))
+        assert flat.shape == (3, 1434)
+
+    def test_predict_1d_promotes(self, stub_emotion_model):
+        assert stub_emotion_model.predict(np.zeros(7)).shape == (1, 1434)
+
+    def test_wrong_width_raises(self, stub_emotion_model):
+        with pytest.raises(ValueError, match="length-7 vector or"):
+            stub_emotion_model.predict(np.zeros((2, 5)))
+
+    def test_scalar_input_raises_clear_error(self, stub_emotion_model):
+        # 0-d / scalar must give a clear ValueError, not a cryptic IndexError
+        with pytest.raises(ValueError, match="length-7 vector"):
+            stub_emotion_model.predict(np.float32(0.5))
+
+    def test_pose_is_implicit_zero(self, stub_emotion_model):
+        # deployed coef has nfeat + 3 rows; predict pads the pose channels to 0
+        out0 = stub_emotion_model.predict(np.zeros(7))
+        assert np.allclose(out0, stub_emotion_model._intercept)
+
+
+class TestPredictFromFeatures:
+    def test_single_returns_478x3(self, stub_emotion_model):
+        mesh = predict_face_mesh_from_features(np.zeros(7), model=stub_emotion_model)
+        assert mesh.shape == (478, 3)
+
+    def test_batch_returns_n478x3(self, stub_blendshape_model):
+        mesh = predict_face_mesh_from_features(np.zeros((5, 52)), model=stub_blendshape_model)
+        assert mesh.shape == (5, 478, 3)
+
+    def test_single_face_2d_row_kept_as_batch(self, stub_emotion_model):
+        # (1, 7) is a batch of one -> (1, 478, 3) from the predict helper
+        mesh = predict_face_mesh_from_features(np.zeros((1, 7)), model=stub_emotion_model)
+        assert mesh.shape == (1, 478, 3)
+
+    def test_axis_major_reshape(self, stub_emotion_model):
+        flat = stub_emotion_model.predict(np.zeros(7))[0]
+        mesh = predict_face_mesh_from_features(np.zeros(7), model=stub_emotion_model)
+        assert np.allclose(mesh[:, 0], flat[:478])
+        assert np.allclose(mesh[:, 1], flat[478:956])
+        assert np.allclose(mesh[:, 2], flat[956:])
+
+    def test_rejects_wrong_model_type(self):
+        with pytest.raises(ValueError, match="PLSFeatMeshModel"):
+            predict_face_mesh_from_features(np.zeros(7), model="not a model")
+
+
+class TestPlotFaceMeshFeatures:
+    def test_plot_emotion(self, stub_emotion_model):
+        emo = np.zeros(7)
+        emo[FEAT_EMOTION_COLUMNS.index("happiness")] = 1.0
+        ax = plot_face_mesh(emotion=emo)
+        assert ax is not None
+
+    def test_plot_blendshapes(self, stub_blendshape_model):
+        bs = np.zeros(52)
+        bs[MP_BLENDSHAPE_NAMES.index("jawOpen")] = 1.0
+        ax = plot_face_mesh(blendshapes=bs)
+        assert ax is not None
+
+    def test_plot_accepts_single_face_2d_row(self, stub_emotion_model):
+        # plot_face_mesh should accept a (1, 7) single-face row, not reject it
+        ax = plot_face_mesh(emotion=np.zeros((1, 7)))
+        assert ax is not None
+
+    def test_mutually_exclusive_inputs(self, stub_emotion_model):
+        with pytest.raises(ValueError, match="at most one"):
+            plot_face_mesh(emotion=np.zeros(7), blendshapes=np.zeros(52))
+
+
+class TestLoaderValidation:
+    def test_bad_feature_name(self):
+        with pytest.raises(ValueError, match="feature must be one of"):
+            plt_mod._load_pls_feat_to_mesh_from_hub("gaze")
+
+
+@pytest.mark.network
+def test_real_emotion_model_from_hub():
+    m = load_emotion_face_mesh_model()
+    assert isinstance(m, PLSFeatMeshModel)
+    assert m.feature_columns == FEAT_EMOTION_COLUMNS
+    mesh = predict_face_mesh_from_features(np.zeros(7), model=m)
+    assert mesh.shape == (478, 3)
+    # happiness should move the mesh away from neutral
+    neu = np.zeros(7)
+    neu[FEAT_EMOTION_COLUMNS.index("neutral")] = 1.0
+    hap = np.zeros(7)
+    hap[FEAT_EMOTION_COLUMNS.index("happiness")] = 1.0
+    d = predict_face_mesh_from_features(hap, model=m) - predict_face_mesh_from_features(neu, model=m)
+    assert np.abs(d).max() > 1.0  # pixel-frame units
+
+
+@pytest.mark.network
+def test_real_blendshape_model_from_hub():
+    m = load_blendshape_face_mesh_model()
+    assert isinstance(m, PLSFeatMeshModel)
+    assert m.feature_columns == MP_BLENDSHAPE_NAMES
+    mesh = predict_face_mesh_from_features(np.zeros(52), model=m)
+    assert mesh.shape == (478, 3)