convert_smpl.py

import argparse
import shutil
import pickle
import numpy as np
import zipfile
import tarfile
from pathlib import Path
from chumpy import Ch

def init_sources():
    """
    Every possible model classification location we may fill.
    """
    return {
        "smpl": {      
            "male": None,
            "female": None,
            "neutral": None,
        },
        "smplx": {      # smplx_locked_head.tar.bz2 and NPZ zip
            "male_pkl": None,
            "male_npz": None,
            "female_pkl": None,
            "female_npz": None,
            "neutral_pkl": None,
            "neutral_npz": None,
            "source": None,    # 'tar' or 'npz-only'
        },
        "smplh": {      
            "male": None,
            "female": None,
            "is_premerged": False,
            "source": None,
        },
        "mano": {      
            "left": None,
            "right": None,
        },
        "merge_candidates": {  # body models from smplh.tar.xz, to merge with MANO
            "male_body": None,
            "female_body": None,
            "neutral_body": None,
        }
    }

def extract_archive(archive_path: Path, extract_root: Path) -> Path:
    """
    Extract zip/tar/tar.bz2 files.
    Returns the directory path where files are extracted.
    """
    target = extract_root / archive_path.stem
    if target.exists():
        return target

    target.mkdir(parents=True, exist_ok=True)
    print(f"[INFO] Extracting {archive_path.name} ...")

    suffix = archive_path.suffix.lower()

    if suffix == ".zip":
        with zipfile.ZipFile(archive_path, 'r') as z:
            z.extractall(target)
    elif suffix in (".gz", ".xz", ".bz2") or archive_path.name.endswith(".tar.bz2") or archive_path.name.endswith(".tar.xz"):
        with tarfile.open(archive_path, 'r:*') as t:
            t.extractall(target)
    else:
        print(f"[WARN] Unsupported archive format: {archive_path.name}\n Skipping extraction.")
        return target

    return target

def detect_models(extract_root: Path, sources: dict):
    """
    Walk all extracted directories and classify paths
    without ever letting smplh.tar overwrite canonical SMPLH.
    """
    for path in extract_root.glob("**/*"):
        if not path.is_file():
            continue

        rel = path.relative_to(extract_root).as_posix().lower()

        # SMPL Base
        if rel.endswith("basicmodel_m_lbs_10_207_0_v1.1.0.pkl"):
            sources["smpl"]["male"] = path
            continue
        if rel.endswith("basicmodel_f_lbs_10_207_0_v1.1.0.pkl"):
            sources["smpl"]["female"] = path
            continue
        if rel.endswith("basicmodel_neutral_lbs_10_207_0_v1.1.0.pkl"):
            sources["smpl"]["neutral"] = path
            continue

        # SMPLX TAR (full)
        if (
            (rel.endswith("model.pkl") or rel.endswith("model.npz"))
            and "smplx_locked_head" in rel
            and any(g in rel for g in ("/male/", "/female/", "/neutral/"))
        ):
            sources["smplx"]["source"] = "tar"

            # assign by gender + extension
            if "/male/" in rel:
                if rel.endswith("model.pkl"): sources["smplx"]["male_pkl"] = path
                if rel.endswith("model.npz"): sources["smplx"]["male_npz"] = path

            elif "/female/" in rel:
                if rel.endswith("model.pkl"): sources["smplx"]["female_pkl"] = path
                if rel.endswith("model.npz"): sources["smplx"]["female_npz"] = path

            elif "/neutral/" in rel:
                if rel.endswith("model.pkl"): sources["smplx"]["neutral_pkl"] = path
                if rel.endswith("model.npz"): sources["smplx"]["neutral_npz"] = path

            continue

        # SMPLX NPZ-only (zip)
        if "models_lockedhead/smplx" in rel and rel.endswith(".npz"):
            if sources["smplx"]["source"] != "tar":
                sources["smplx"]["source"] = "npz-only"
            fname = path.name.lower()
            if "male" in fname: sources["smplx"]["male_npz"] = path
            elif "female" in fname: sources["smplx"]["female_npz"] = path
            elif "neutral" in fname: sources["smplx"]["neutral_npz"] = path
            continue

        # SMPLH (ONLY FROM smplx.zip)
        if "smplx/smplh" in rel and path.name.startswith("SMPLH_"):
            sources["smplh"]["source"] = "premerged"
            sources["smplh"]["is_premerged"] = True
            if "smplh_male" in path.name.lower():
                sources["smplh"]["male"] = path
            elif "smplh_female" in path.name.lower():
                sources["smplh"]["female"] = path
            continue

        # SMPLH TAR BODY MODELS (for MERGE ONLY)
        if "smplh" in rel and rel.endswith("model.npz"):
            gender = None
            if "/male/" in rel: gender = "male_body"
            elif "/female/" in rel: gender = "female_body"
            elif "/neutral/" in rel: gender = "neutral_body"
            if gender:
                sources["merge_candidates"][gender] = path
            continue

        # MANO
        if "mano_v1_2/models/mano_left.pkl" in rel:
            sources["mano"]["left"] = path
            continue
        if "mano_v1_2/models/mano_right.pkl" in rel:
            sources["mano"]["right"] = path
            continue

    # Resolve SMPLX source
    smplx = sources["smplx"]
    if any([
        smplx["male_pkl"], smplx["female_pkl"], smplx["neutral_pkl"],
    ]):
        smplx["source"] = "tar"
    elif any([smplx["male_npz"], smplx["female_npz"], smplx["neutral_npz"]]):
        if not smplx["source"]:
            smplx["source"] = "npz-only"
    else:
        smplx["source"] = None

    return sources

def load_model(path):
    """Load pickle or NPZ uniformly."""
    ext = path.suffix.lower()
    if ext == ".npz":
        npz = np.load(path, allow_pickle=True)
        return {k: npz[k] for k in npz.files}
    with open(path, "rb") as f:
        return pickle.load(f, encoding="latin1")

def sanitize(data, convert_csc_matrix=False):
    """Clean SMPL-family metadata."""
    out = {}
    for k, v in data.items():
        if isinstance(v, Ch):
            v = v.r 
        
        if convert_csc_matrix and "csc_matrix" in str(type(v)):
            out[k] = v.toarray().astype(np.float32)
            continue
        if isinstance(v, np.ndarray):
            if np.issubdtype(v.dtype, np.number):
                if v.dtype.kind == "f":
                    out[k] = v.astype(np.float32)
                else:
                    out[k] = v.astype(np.int32)
            else:
                out[k] = v
            continue
        if isinstance(v, (float, np.floating)):
            out[k] = np.float32(v)
            continue
        if isinstance(v, (int, np.integer)):
            out[k] = np.int32(v)
            continue
        if isinstance(v, list) and all(isinstance(x,(int,float)) for x in v):
            out[k] = np.array(v, dtype=np.float32)
            continue
        out[k] = v
    return out

def save_model(obj, outpath, save_npz):
    outpath.parent.mkdir(parents=True, exist_ok=True)
    if save_npz:
        np.savez_compressed(outpath, **obj)
    else:
        with open(outpath, "wb") as f:
            pickle.dump(obj, f)
    print(f"  ✓ Saved → {outpath}")
    
def merge_mano_smplh(smplh_body_path, mano_left, mano_right, out_path):
    """Merge MANO hands into SMPLH body model and save."""
    print(f"   → Merging MANO+SMPLH: {out_path.name}")

    body = load_model(smplh_body_path)
    left = load_model(mano_left)
    right = load_model(mano_right)

    # drop any old hand keys
    for k in list(body.keys()):
        if k.startswith(("hands_components","hands_coeffs","hands_mean")):
            del body[k]

    # merge new ones
    body["hands_componentsl"] = left["hands_components"]
    body["hands_componentsr"] = right["hands_components"]
    body["hands_coeffsl"] = left["hands_coeffs"]
    body["hands_coeffsr"] = right["hands_coeffs"]
    body["hands_meanl"] = left["hands_mean"]
    body["hands_meanr"] = right["hands_mean"]

    save_model(body, out_path, save_npz=False)  

def write_smpl_models(sources, out_root, args):
    """Write canonical SMPL if found."""
    smpl = sources["smpl"]
    for gender, path in smpl.items():
        if gender == "all": continue
        if path is None: continue

        data = sanitize(load_model(path), args.convert_csc_matrix)
        suffix = ".npz" if args.save_as_npz else ".pkl"
        fname = f"SMPL_{gender.upper()}{suffix}"
        save_model(data, out_root/"smpl"/fname, args.save_as_npz)

def write_smplx_models(sources, out_root, args):
    """Write canonical + alternate SMPLX."""
    smplx = sources["smplx"]
    main_dir = out_root/"smplx"

    def write(path, tag):
        if path:
            d = sanitize(load_model(path), args.convert_csc_matrix)
            suffix = ".npz" if path.suffix.lower()==".npz" and args.save_as_npz else ".pkl"
            out = main_dir/f"SMPLX_{tag.upper()}{suffix}"
            save_model(d, out, args.save_as_npz)

    src = smplx["source"]
    if src=="tar":
        print("[INFO] SMPLX source: TAR (locked_head.tar.bz2)")
        write(smplx["male_pkl"], "male")
        write(smplx["male_npz"], "male")
        write(smplx["female_pkl"], "female")
        write(smplx["female_npz"], "female")
        write(smplx["neutral_pkl"], "neutral")
        write(smplx["neutral_npz"], "neutral")

    elif src=="npz-only":
        print("[INFO] SMPLX source: NPZ-only zip")
        for key in ["male_npz","female_npz","neutral_npz"]:
            path = smplx[key]
            if path:
                d = sanitize(load_model(path), args.convert_csc_matrix)
                suffix = ".npz" if args.save_as_npz else ".pkl"
                out = main_dir/f"SMPLX_{key.split('_')[0].upper()}{suffix}"
                save_model(d, out, args.save_as_npz)
    else:
        print("[WARN] No SMPLX models detected. Check you downloaded everything correctly.")

def write_smplh_models(sources, out_root, args):
    """Write canonical SMPLH from smplx.zip."""
    smplh = sources["smplh"]
    main_dir = out_root/"smplh"

    if smplh["source"] == "premerged":
        print("[INFO] SMPLH source: pre-merged SMPLH from smplx.zip")

        for gender in ["male","female"]:
            path = smplh[gender]
            if path:
                d = sanitize(load_model(path), args.convert_csc_matrix)
                suffix = ".npz" if args.save_as_npz else ".pkl"
                save_model(d, main_dir/f"SMPLH_{gender.upper()}{suffix}", args.save_as_npz)
    else:
        print("[INFO] SMPLH (smplx.zip) found. Check you downloaded everything correctly.")

def write_mano(sources, out_root, args):
    """Save only MANO_LEFT / MANO_RIGHT."""
    mano = sources["mano"]
    main_dir = out_root/"mano"

    for tag in ["left","right"]:
        path = mano[tag]
        if path:
            d = sanitize(load_model(path), args.convert_csc_matrix)
            suffix = ".npz" if args.save_as_npz else ".pkl"
            name = f"MANO_{tag.upper()}{suffix}"
            save_model(d, main_dir/name, args.save_as_npz)

def perform_merge(sources, out_root):
    """Merge SMPLH body from smplh.tar.xz + MANO from mano_v1_2."""

    body = sources["merge_candidates"]
    mano = sources["mano"]

    if not any([body["male_body"], body["female_body"], body["neutral_body"]]):
        print("[WARN] No SMPLH body models found for merging (smplh.tar.xz). Check you downloaded everything correctly.")
        return

    if not (mano["left"] and mano["right"]):
        print("[WARN] Missing MANO_LEFT and/or MANO_RIGHT, cannot merge.")
        return

    print("\n=== MERGING SMPLH + MANO ===")

    merged_dir = out_root / "smplh_merged"
    merged_dir.mkdir(parents=True, exist_ok=True)

    mapping = {
        "male_body":    "male",
        "female_body":  "female",
        "neutral_body": "neutral",
    }

    for key, body_path in body.items():
        if body_path is None:
            continue
        gender = mapping[key]  
        out_path = merged_dir / f"SMPLH_{gender.upper()}.pkl"
        merge_mano_smplh(
            smplh_body_path=body_path,
            mano_left=mano["left"],
            mano_right=mano["right"],
            out_path=out_path
        )
        
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--zips_path", type=str, default="zips",
                        help="Folder containing downloaded model zips/tars")
    parser.add_argument("--output_path", type=str, default="models")
    parser.add_argument("--convert_csc_matrix", action="store_true")
    parser.add_argument("--save_as_npz", action="store_true")
    parser.add_argument("--merge_mano", action="store_true")

    args = parser.parse_args()

    zips_dir = Path(args.zips_path)
    out_root = Path(args.output_path)
    extract_root = zips_dir/"_unpacked"

    if extract_root.exists():
        shutil.rmtree(extract_root)
        extract_root.mkdir(parents=True, exist_ok=True)
    else:
        extract_root.mkdir(parents=True, exist_ok=True)


    print("\n=== SMPL FAMILY CONVERSION ===\n")

    for arch in zips_dir.iterdir():
        if arch.is_file():
            extract_archive(arch, extract_root)

    sources = init_sources()
    detect_models(extract_root, sources)

    print("\n=== DETECTED SOURCES ===")
    for k, v in sources.items():
        print(f"[{k}]")
        for kk, vv in v.items():
            if kk=="all": continue
            print(f"  {kk}: {vv}")
        print()

    write_smpl_models(sources, out_root, args)
    write_smplx_models(sources, out_root, args)
    write_smplh_models(sources, out_root, args)
    write_mano(sources, out_root, args)

    # Perform merge if requested
    if args.merge_mano:
        perform_merge(sources, out_root)

    print("\n=== DONE ===")
    print(f"Models written to: {out_root}")