diff --git a/feat/plotting.py b/feat/plotting.py index 6a3871c9..609ab061 100644 --- a/feat/plotting.py +++ b/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() diff --git a/feat/tests/test_feat_mesh_viz.py b/feat/tests/test_feat_mesh_viz.py new file mode 100644 index 00000000..1543c97d --- /dev/null +++ b/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) diff --git a/feat/tests/test_landmarks68_to_mesh478.py b/feat/tests/test_landmarks68_to_mesh478.py index 8ba1c755..47c9e57e 100644 --- a/feat/tests/test_landmarks68_to_mesh478.py +++ b/feat/tests/test_landmarks68_to_mesh478.py @@ -172,7 +172,7 @@ def test_accepts_precomputed_mesh(self, stub_bridge_model): def test_rejects_both_au_and_mesh(self, stub_bridge_model): mesh = stub_bridge_model.mean_predicted_mesh - with pytest.raises(ValueError, match=r"either `au` or `mesh`"): + with pytest.raises(ValueError, match=r"at most one of"): plot_face_mesh(au=np.zeros(20), mesh=mesh) def test_rejects_wrong_mesh_shape(self):