208 lines
6.1 KiB
Python
208 lines
6.1 KiB
Python
"""
|
|
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),
|
|
)
|
|
}
|