test.py

import argparse
from pathlib import Path
import pickle
import numpy as np
import sys
import traceback

# ============================================================
# Stage 1 — structural check (no torch/smplx required)
#
# Just confirms every .pkl/.npz under --models_path loads and is a
# well-formed dict of numeric arrays. Cheap, but doesn't prove the file
# is actually usable by the smplx library — see Stage 2 below.
# ============================================================

def load_model_file(path: Path):
    """Load pkl or npz model file and return a dict."""
    if path.suffix.lower() == ".npz":
        try:
            npz = np.load(path, allow_pickle=True)
            return {k: npz[k] for k in npz.files}
        except Exception as e:
            raise RuntimeError(f"Failed loading NPZ file: {e}")

    elif path.suffix.lower() == ".pkl":
        try:
            with open(path, "rb") as f:
                return pickle.load(f, encoding="latin1")
        except Exception as e:
            raise RuntimeError(f"Failed loading PKL file: {e}")

    else:
        raise RuntimeError(f"Unsupported model format: {path.name}")

def validate_model(model_dict, path: Path):
    """
    Basic structural checks:
      - must be a dict
      - values should be numeric arrays or common SMPL data types
    """
    if not isinstance(model_dict, dict):
        raise ValueError(f"{path.name} → Loaded object is NOT a dict!")

    count = 0
    types = {}
    for k, v in model_dict.items():
        if isinstance(v, np.ndarray):
            continue
        if isinstance(v, (int, float, np.integer, np.floating)):
            continue
        if isinstance(v, list) and all(isinstance(x, (int, float)) for x in v):
            continue
        if isinstance(v, (dict, str, bytes, type(None))):
            continue
        if 'csc_matrix' in type(v).__name__:
            continue

        t = type(v).__name__
        types[t] = types.get(t, 0) + 1
        count += 1

    print(f"    → {count} unrecognized entries with types: {types}")

def run_structural_checks(root: Path):
    model_files = []
    for ext in ("*.pkl", "*.npz"):
        model_files.extend(root.rglob(ext))

    if not model_files:
        print("[WARN] No .pkl or .npz model files found.")
        return [], []

    passed, failed = [], []
    for f in sorted(model_files):
        label = str(f.relative_to(root))
        print(f"→ Testing: {label}")
        try:
            data = load_model_file(f)
            validate_model(data, f)
        except Exception:
            print(f"  ✗ ERROR loading {f.name}")
            failed.append((label, traceback.format_exc()))
            continue
        print("  ✓ OK")
        passed.append(label)

    return passed, failed

# ============================================================
# Stage 2 — functional check via the real smplx library
#
# Loads every model exactly the way downstream code would, using
# smplx.create() (https://github.com/vchoutas/smplx), and runs a forward
# pass. This is what actually catches things like a model missing hand
# PCA components — structural checks alone can't, since a body-only file
# is just as "well-formed" as a full hand-articulated one.
# ============================================================

# smplh/amass_16betas is a non-standard subfolder (merged SMPL+H + MANO,
# kept separate from smplh/ because the shape space differs), so it can't
# go through smplx.create()'s model_path/model_type resolution — it's
# loaded directly via the SMPLH class instead.
AMASS_MERGED_SUBDIR = Path("smplh") / "amass_16betas"

def discover_functional_targets(root: Path):
    """Find every model instance smplx.create() should be able to load,
    based on the standard body_models/<type>/<TYPE>_<GENDER>.<ext> layout.
    """
    targets = []

    for model_type in ("smpl", "smplx", "smplh"):
        d = root / model_type
        if not d.is_dir():
            continue
        for f in sorted(d.glob(f"{model_type.upper()}_*")):
            ext = f.suffix.lower().lstrip(".")
            if ext not in ("pkl", "npz"):
                continue
            # smplx's SMPL class always reads .pkl regardless of what's
            # passed, so testing it against .npz would be meaningless.
            if model_type == "smpl" and ext != "pkl":
                continue
            gender = f.stem.split("_", 1)[1].lower()
            targets.append({
                "label": f"{model_type}/{f.name}",
                "model_type": model_type,
                "model_path": str(root),
                "kwargs": {"gender": gender, "ext": ext},
                "hand_capable": model_type in ("smplx", "smplh"),
            })

    mano_dir = root / "mano"
    if mano_dir.is_dir():
        for f in sorted(mano_dir.glob("MANO_*")):
            ext = f.suffix.lower().lstrip(".")
            if ext not in ("pkl", "npz"):
                continue
            side = f.stem.split("_", 1)[1].lower()
            targets.append({
                "label": f"mano/{f.name}",
                "model_type": "mano",
                "model_path": str(root),
                "kwargs": {"is_rhand": side == "right", "ext": ext},
                "hand_capable": True,
            })

    amass_dir = root / AMASS_MERGED_SUBDIR
    if amass_dir.is_dir():
        for f in sorted(amass_dir.glob("SMPLH_*")):
            ext = f.suffix.lower().lstrip(".")
            if ext not in ("pkl", "npz"):
                continue
            gender = f.stem.split("_", 1)[1].lower()
            targets.append({
                "label": f"smplh/amass_16betas/{f.name}",
                "model_type": "smplh",
                "model_path": str(amass_dir),
                "kwargs": {"gender": gender, "ext": ext},
                "hand_capable": True,
            })

    return targets

def instantiate(smplx, target, use_pca):
    kwargs = dict(target["kwargs"])
    if target["hand_capable"]:
        kwargs["use_pca"] = use_pca
    if target["model_path"].endswith(str(AMASS_MERGED_SUBDIR)):
        return smplx.SMPLH(target["model_path"], batch_size=1, **kwargs)
    return smplx.create(target["model_path"], model_type=target["model_type"],
                         batch_size=1, **kwargs)

def discover_amass_native_targets(root: Path):
    """Find smplh/<gender>/model.npz files — the raw, un-merged AMASS body
    used directly by human_body_prior.BodyModel (not the smplx package)."""
    targets = []
    smplh_dir = root / "smplh"
    if not smplh_dir.is_dir():
        return targets
    for gender_dir in sorted(smplh_dir.iterdir()):
        if not gender_dir.is_dir() or gender_dir.name not in ("male", "female", "neutral"):
            continue
        f = gender_dir / "model.npz"
        if f.is_file():
            targets.append({"label": f"smplh/{gender_dir.name}/model.npz", "path": f})
    return targets

def check_amass_native_smplh(torch, smplx, path: Path):
    """Functionally validate a raw AMASS-native SMPLH file the way
    human_body_prior.BodyModel would use it — no MANO merge, hand pose is
    raw axis-angle, and it never densifies sparse matrices itself, so a
    stray scipy sparse J_regressor (usually caught by smplx's own to_np()
    for the other model types) would silently break it here.

    human_body_prior isn't installed as a test dependency (heavy, license
    file, unrelated extras), so this reimplements just enough of its
    forward pass using smplx's own lbs() — the same skinning math — to
    prove the file actually deforms correctly.
    """
    npz = np.load(path, allow_pickle=True)
    required = ("v_template", "f", "shapedirs", "J_regressor", "posedirs", "kintree_table", "weights")
    missing = [k for k in required if k not in npz.files]
    if missing:
        raise ValueError(f"missing keys required by human_body_prior.BodyModel: {missing}")

    for k in ("shapedirs", "J_regressor", "posedirs", "weights"):
        if npz[k].dtype == object:
            raise ValueError(f"'{k}' isn't a plain dense array — BodyModel loads it with a bare "
                              f"torch.tensor(...) and has no sparse-matrix handling")

    v_template = torch.tensor(npz["v_template"], dtype=torch.float32)
    shapedirs = torch.tensor(npz["shapedirs"], dtype=torch.float32)
    posedirs_raw = torch.tensor(npz["posedirs"], dtype=torch.float32)
    J_regressor = torch.tensor(npz["J_regressor"], dtype=torch.float32)
    kintree_table = npz["kintree_table"].astype(np.int64)
    weights = torch.tensor(npz["weights"], dtype=torch.float32)
    faces = npz["f"]

    n_verts = v_template.shape[0]
    n_joints = kintree_table.shape[1]
    n_betas = shapedirs.shape[-1]

    # Same reshape human_body_prior.BodyModel.__init__ applies before lbs().
    posedirs = posedirs_raw.reshape(n_verts * 3, -1).T

    betas = torch.zeros(1, n_betas)
    pose = torch.zeros(1, n_joints * 3)
    parents = torch.tensor(kintree_table[0], dtype=torch.long)

    verts, joints = smplx.lbs.lbs(
        betas=betas, pose=pose,
        v_template=v_template.unsqueeze(0),
        shapedirs=shapedirs, posedirs=posedirs,
        J_regressor=J_regressor, parents=parents,
        lbs_weights=weights,
    )

    if not torch.isfinite(verts).all():
        raise ValueError("model produced non-finite vertices")
    max_face_idx = int(faces.max())
    if max_face_idx >= n_verts:
        raise ValueError(f"faces reference vertex {max_face_idx} but model only has {n_verts} vertices")

    return n_verts, n_joints, n_betas

def run_functional_checks(root: Path):
    try:
        import torch
        import smplx
    except ImportError:
        print("[WARN] torch/smplx not installed — skipping functional checks.")
        print("       Install with: pip install torch smplx")
        return None, None

    targets = discover_functional_targets(root)
    amass_native_targets = discover_amass_native_targets(root)
    if not targets and not amass_native_targets:
        print("[WARN] No models found matching the expected body_models/ layout.")
        return [], []

    passed, failed = [], []
    for target in targets:
        # Hand-capable models are tested with PCA-compressed hand pose
        # (the common case) and with the raw 45-DoF pose per hand
        # (use_pca=False) — the exact path that breaks if hand PCA
        # components are missing from the file.
        pca_variants = (True, False) if target["hand_capable"] else (True,)
        for use_pca in pca_variants:
            suffix = ""
            if target["hand_capable"]:
                suffix = f" (use_pca={use_pca})"
            label = target["label"] + suffix
            try:
                model = instantiate(smplx, target, use_pca)
                with torch.no_grad():
                    output = model()
                verts = output.vertices
                joints = output.joints
                if not torch.isfinite(verts).all():
                    raise ValueError("model produced non-finite vertices")
                if not torch.isfinite(joints).all():
                    raise ValueError("model produced non-finite joints")
                n_verts = verts.shape[1]
                max_face_idx = int(model.faces_tensor.max().item())
                if max_face_idx >= n_verts:
                    raise ValueError(
                        f"faces reference vertex {max_face_idx} but model "
                        f"only has {n_verts} vertices"
                    )
                print(f"  ✓ {label}: {n_verts} verts, {joints.shape[1]} joints")
                passed.append(label)
            except Exception as e:
                print(f"  ✗ {label}: {e}")
                failed.append((label, traceback.format_exc()))

    for target in amass_native_targets:
        label = target["label"] + " (human_body_prior-native)"
        try:
            n_verts, n_joints, n_betas = check_amass_native_smplh(torch, smplx, target["path"])
            print(f"  ✓ {label}: {n_verts} verts, {n_joints} joints, {n_betas} betas")
            passed.append(label)
        except Exception as e:
            print(f"  ✗ {label}: {e}")
            failed.append((label, traceback.format_exc()))

    return passed, failed

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--models_path",
        type=str,
        default='body_models',
        help="Path to the directory containing the converted models.",
    )
    parser.add_argument(
        "--skip-functional",
        action="store_true",
        help="Only run the structural (dict-of-arrays) check; skip loading "
             "models through the smplx library.",
    )
    args = parser.parse_args()

    root = Path(args.models_path)
    if not root.exists():
        print(f"[ERROR] Path does not exist: {root}")
        sys.exit(1)

    print(f"\n=== STAGE 1: STRUCTURAL CHECK — {root} ===\n")
    struct_passed, struct_failed = run_structural_checks(root)

    func_passed, func_failed = None, None
    if not args.skip_functional:
        print(f"\n=== STAGE 2: FUNCTIONAL CHECK (smplx.create) — {root} ===\n")
        func_passed, func_failed = run_functional_checks(root)

    print("\n================================")
    print("           SUMMARY")
    print("================================\n")

    print(f"Structural: {len(struct_passed)} passed, {len(struct_failed)} failed")
    if func_failed is None:
        print("Functional: skipped")
    else:
        print(f"Functional: {len(func_passed)} passed, {len(func_failed)} failed")

    any_failed = bool(struct_failed) or bool(func_failed)

    if struct_failed:
        print("\n--- STRUCTURAL FAILURES ---")
        for label, msg in struct_failed:
            print(f"\n{label}\n{msg}")

    if func_failed:
        print("\n--- FUNCTIONAL FAILURES ---")
        for label, msg in func_failed:
            print(f"\n{label}\n{msg}")

    print("\n=== DONE ===")
    sys.exit(1 if any_failed else 0)