Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+207
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
Utilities for loading input and models (reading pre-processed data)
|
||||
"""
|
||||
|
||||
import os
|
||||
import importlib
|
||||
import glob
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# from body_model import HandModel
|
||||
|
||||
|
||||
ROOT_DIR = os.path.abspath(f"{__file__}/../../../")
|
||||
|
||||
|
||||
def load_config_from_log(log_dir):
|
||||
hydra_dir = f"{log_dir}/.hydra"
|
||||
cfg_path = f"{hydra_dir}/config.yaml"
|
||||
assert os.path.isdir(hydra_dir), f"{hydra_dir} does not exist"
|
||||
assert os.path.isfile(cfg_path), f"{cfg_path} does not exist"
|
||||
return OmegaConf.load(cfg_path)
|
||||
|
||||
|
||||
def resolve_cfg_paths(cfg):
|
||||
paths = cfg.paths
|
||||
for name, rel_path in paths.items():
|
||||
if rel_path.startswith("/"): # absolute path
|
||||
continue
|
||||
paths[name] = os.path.join(ROOT_DIR, rel_path)
|
||||
print("RESOLVED PATHS", paths)
|
||||
return cfg
|
||||
|
||||
|
||||
# def load_mano_body_model(cfg):
|
||||
# """
|
||||
# Load SMPL model
|
||||
# """
|
||||
# if device is None:
|
||||
# device = torch.device("cpu")
|
||||
# return (
|
||||
# HandModel(
|
||||
# bm_path=path,
|
||||
# num_betas=num_betas,
|
||||
# batch_size=batch_size,
|
||||
# use_vtx_selector=use_vtx_selector,
|
||||
# model_type=model_type,
|
||||
# ).to(device),
|
||||
# )
|
||||
|
||||
|
||||
def expid2model(expr_dir):
|
||||
""" "
|
||||
Reading VPoser models (https://github.com/nghorbani/human_body_prior).
|
||||
"""
|
||||
from configer import Configer
|
||||
|
||||
if not os.path.exists(expr_dir):
|
||||
raise ValueError("Could not find the experiment directory: %s" % expr_dir)
|
||||
|
||||
best_model_fname = sorted(
|
||||
glob.glob(os.path.join(expr_dir, "snapshots", "*.pt")), key=os.path.getmtime
|
||||
)[-1]
|
||||
try_num = os.path.basename(best_model_fname).split("_")[0]
|
||||
|
||||
print(("Found Trained Model: %s" % best_model_fname))
|
||||
|
||||
default_ps_fname = glob.glob(os.path.join(expr_dir, "*.ini"))[0]
|
||||
if not os.path.exists(default_ps_fname):
|
||||
raise ValueError(
|
||||
"Could not find the appropriate vposer_settings: %s" % default_ps_fname
|
||||
)
|
||||
ps = Configer(
|
||||
default_ps_fname=default_ps_fname,
|
||||
work_dir=expr_dir,
|
||||
best_model_fname=best_model_fname,
|
||||
)
|
||||
|
||||
return ps, best_model_fname
|
||||
|
||||
|
||||
def load_vposer(expr_dir, vp_model="snapshot"):
|
||||
"""
|
||||
:param expr_dir:
|
||||
:param vp_model: either 'snapshot' to use the experiment folder's code or a VPoser imported module, e.g.
|
||||
from human_body_prior.train.vposer_smpl import VPoser, then pass VPoser to this function
|
||||
:param if True will load the model definition used for training, and not the one in current repository
|
||||
:return:
|
||||
"""
|
||||
ps, trained_model_fname = expid2model(expr_dir)
|
||||
if vp_model == "snapshot":
|
||||
|
||||
vposer_path = sorted(
|
||||
glob.glob(os.path.join(expr_dir, "vposer_*.py")), key=os.path.getmtime
|
||||
)[-1]
|
||||
|
||||
spec = importlib.util.spec_from_file_location("VPoser", vposer_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
vposer_pt = getattr(module, "VPoser")(
|
||||
num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape=ps.data_shape
|
||||
)
|
||||
else:
|
||||
vposer_pt = vp_model(
|
||||
num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape=ps.data_shape
|
||||
)
|
||||
|
||||
vposer_pt.load_state_dict(torch.load(trained_model_fname, map_location="cpu"))
|
||||
vposer_pt.eval()
|
||||
|
||||
return vposer_pt, ps
|
||||
|
||||
|
||||
def load_state(
|
||||
load_path,
|
||||
model,
|
||||
optimizer=None,
|
||||
is_parallel=False,
|
||||
map_location=None,
|
||||
ignore_keys=None,
|
||||
):
|
||||
"""
|
||||
Load Humor model checkpoint
|
||||
"""
|
||||
if not os.path.exists(load_path):
|
||||
print("Could not find checkpoint at path " + load_path)
|
||||
|
||||
full_checkpoint_dict = torch.load(load_path, map_location=map_location)
|
||||
model_state_dict = full_checkpoint_dict["model"]
|
||||
optim_state_dict = full_checkpoint_dict["optim"]
|
||||
|
||||
# load model weights
|
||||
for k, v in model_state_dict.items():
|
||||
if k.split(".")[0] == "module" and not is_parallel:
|
||||
# then it was trained with Data parallel
|
||||
print("Loading weights trained with DataParallel...")
|
||||
model_state_dict = {
|
||||
".".join(k.split(".")[1:]): v
|
||||
for k, v in model_state_dict.items()
|
||||
if k.split(".")[0] == "module"
|
||||
}
|
||||
break
|
||||
|
||||
if ignore_keys is not None:
|
||||
model_state_dict = {
|
||||
k: v
|
||||
for k, v in model_state_dict.items()
|
||||
if k.split(".")[0] not in ignore_keys
|
||||
}
|
||||
|
||||
# overwrite entries in the existing state dict
|
||||
missing_keys, unexpected_keys = model.load_state_dict(
|
||||
model_state_dict, strict=False
|
||||
)
|
||||
if ignore_keys is not None:
|
||||
missing_keys = [k for k in missing_keys if k.split(".")[0] not in ignore_keys]
|
||||
unexpected_keys = [
|
||||
k for k in unexpected_keys if k.split(".")[0] not in ignore_keys
|
||||
]
|
||||
if len(missing_keys) > 0:
|
||||
print(
|
||||
"WARNING: The following keys could not be found in the given state dict - ignoring..."
|
||||
)
|
||||
print(missing_keys)
|
||||
if len(unexpected_keys) > 0:
|
||||
print(
|
||||
"WARNING: The following keys were found in the given state dict but not in the current model - ignoring..."
|
||||
)
|
||||
print(unexpected_keys)
|
||||
|
||||
# load optimizer weights
|
||||
if optimizer is not None:
|
||||
optimizer.load_state_dict(optim_state_dict)
|
||||
|
||||
min_train_loss = float("Inf")
|
||||
if "min_train_loss" in full_checkpoint_dict.keys():
|
||||
min_train_loss = full_checkpoint_dict["min_train_loss"]
|
||||
|
||||
return (
|
||||
full_checkpoint_dict["epoch"],
|
||||
full_checkpoint_dict["min_val_loss"],
|
||||
min_train_loss,
|
||||
)
|
||||
|
||||
|
||||
def load_gmm(path, device=None):
|
||||
"""
|
||||
Load motion prior GMM for first frame motion prior
|
||||
"""
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
|
||||
gmm_path = os.path.join(path, "prior_gmm.npz")
|
||||
if not os.path.isfile(gmm_path):
|
||||
raise ValueError(f"Init motion prior path {gmm_path} does not exist")
|
||||
|
||||
res = np.load(gmm_path)
|
||||
return {
|
||||
"gmm": (
|
||||
torch.from_numpy(res["weights"]).to(device),
|
||||
torch.from_numpy(res["means"]).to(device),
|
||||
torch.from_numpy(res["covariances"]).to(device),
|
||||
)
|
||||
}
|
||||
Reference in New Issue
Block a user