Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
from . import loaders
|
||||
from . import logger
|
||||
from . import tensor
|
||||
+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),
|
||||
)
|
||||
}
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
import datetime
|
||||
|
||||
|
||||
class Logger(object):
|
||||
"""
|
||||
Static logging class
|
||||
"""
|
||||
|
||||
log_file = None
|
||||
|
||||
@staticmethod
|
||||
def init(log_path):
|
||||
Logger.log_file = log_path
|
||||
|
||||
@staticmethod
|
||||
def log(write_str, to_stdout=True):
|
||||
if to_stdout:
|
||||
print(write_str)
|
||||
if not Logger.log_file:
|
||||
print("Logger must be initialized before logging!")
|
||||
return
|
||||
time_str = datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S")
|
||||
with open(Logger.log_file, "a") as f:
|
||||
f.write(time_str + " ")
|
||||
f.write(str(write_str) + "\n")
|
||||
|
||||
|
||||
def log_cur_stats(stats_dict, iter=None, to_stdout=True):
|
||||
loss = stats_dict.pop("total", 0)
|
||||
Logger.log(f"LOSS: {loss:.04f}", to_stdout=to_stdout)
|
||||
for k, v in stats_dict.items():
|
||||
Logger.log(f"{k}: {v:.04f}", to_stdout=to_stdout)
|
||||
if to_stdout:
|
||||
if iter is not None:
|
||||
print("======= iter %d =======" % iter)
|
||||
else:
|
||||
print("========")
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def move_to(obj, device):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.to(device)
|
||||
if isinstance(obj, dict):
|
||||
return {k: move_to(v, device) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [move_to(x, device) for x in obj]
|
||||
return obj # otherwise do nothing
|
||||
|
||||
|
||||
def detach_all(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.detach()
|
||||
if isinstance(obj, dict):
|
||||
return {k: detach_all(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [detach_all(x) for x in obj]
|
||||
return obj # otherwise do nothing
|
||||
|
||||
|
||||
def to_torch(obj):
|
||||
if isinstance(obj, np.ndarray):
|
||||
return torch.from_numpy(obj).float()
|
||||
if isinstance(obj, dict):
|
||||
return {k: to_torch(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [to_torch(x) for x in obj]
|
||||
return obj
|
||||
|
||||
|
||||
def to_np(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.numpy()
|
||||
if isinstance(obj, dict):
|
||||
return {k: to_np(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [to_np(x) for x in obj]
|
||||
return obj
|
||||
|
||||
|
||||
def get_device(i=0):
|
||||
device = f"cuda:{i}" if torch.cuda.is_available() else "cpu"
|
||||
return torch.device(device)
|
||||
|
||||
|
||||
def scatter_intervals(tensor, start, end, T=None):
|
||||
"""
|
||||
Scatter the tensor contents into intervals from start to end
|
||||
output tensor indexed from 0 to end.max()
|
||||
:param tensor (B, S, *)
|
||||
:param start (B) start indices
|
||||
:param end (B) end indices
|
||||
:param T (int, optional) max length
|
||||
returns (B, T, *) scattered tensor
|
||||
"""
|
||||
assert isinstance(tensor, torch.Tensor) and tensor.ndim >= 2
|
||||
if T is not None:
|
||||
end = torch.where(end < 0, T + 1 + end, end)
|
||||
assert torch.all(end <= T)
|
||||
else:
|
||||
T = end.max()
|
||||
|
||||
B, S, *dims = tensor.shape
|
||||
start, end = start.long(), end.long()
|
||||
# get idcs that go past the last time step so we don't have repeat indices in scatter
|
||||
idcs = time_segment_idcs(start, end, clip=False) # (B, S)
|
||||
# mask out the extra padding
|
||||
mask = idcs >= end[:, None]
|
||||
tensor[mask] = 0
|
||||
|
||||
idcs = idcs.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
|
||||
output = torch.zeros(
|
||||
B, idcs.max() + 1, *dims, device=tensor.device, dtype=tensor.dtype
|
||||
)
|
||||
output.scatter_(1, idcs, tensor)
|
||||
# slice out the extra segments
|
||||
return output[:, :T]
|
||||
|
||||
|
||||
def get_scatter_mask(start, end, T):
|
||||
"""
|
||||
get the mask of selected intervals
|
||||
"""
|
||||
B = start.shape[0]
|
||||
start, end = start.long(), end.long()
|
||||
assert torch.all(end <= T)
|
||||
idcs = time_segment_idcs(start, end, clip=True)
|
||||
mask = torch.zeros(B, T, device=start.device, dtype=torch.bool)
|
||||
mask.scatter_(1, idcs, 1)
|
||||
return mask
|
||||
|
||||
|
||||
def select_intervals(series, start, end):
|
||||
"""
|
||||
Select slices of a tensor from start to end
|
||||
will pad uneven sequences to all the max segment length
|
||||
:param series (B, T, *)
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
returns (B, S, *) selected segments, S = max(end - start)
|
||||
"""
|
||||
B, T, *dims = series.shape
|
||||
assert torch.all(end <= T)
|
||||
sel = time_segment_idcs(start, end, clip=True)
|
||||
S = sel.shape[1]
|
||||
sel = sel.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
|
||||
return torch.gather(series, 1, sel)
|
||||
|
||||
|
||||
def get_select_mask(start, end):
|
||||
"""
|
||||
get the mask of unpadded elementes for the selected time segments
|
||||
e.g. sel[mask] are the unpadded elements
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
"""
|
||||
idcs = time_segment_idcs(start, end, clip=False)
|
||||
return idcs < end[:, None] # (B, S)
|
||||
|
||||
|
||||
def time_segment_idcs(start, end, clip=True):
|
||||
"""
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
returns (B, S) long tensor of indices, where S = max(end - start)
|
||||
"""
|
||||
start, end = start.long(), end.long()
|
||||
S = (end - start).max()
|
||||
seg = torch.arange(S, dtype=torch.int64, device=start.device)
|
||||
idcs = start[:, None] + seg[None, :] # (B, S)
|
||||
if clip:
|
||||
# clip at the lengths of each track
|
||||
imax = end[:, None] - 1
|
||||
idcs = idcs.clamp(max=imax)
|
||||
return idcs
|
||||
Reference in New Issue
Block a user