Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+204
@@ -0,0 +1,204 @@
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import subprocess
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from data import get_dataset_from_cfg, expand_source_paths
|
||||
|
||||
# from humor.humor_model import HumorModel
|
||||
from optim.base_scene import BaseSceneModel
|
||||
|
||||
from optim.optimizers import (
|
||||
RootOptimizer,
|
||||
SmoothOptimizer
|
||||
)
|
||||
from optim.output import (
|
||||
save_track_info,
|
||||
save_camera_json,
|
||||
save_input_poses,
|
||||
save_initial_predictions,
|
||||
)
|
||||
from vis.viewer import init_viewer
|
||||
from body_model import MANO
|
||||
from util.loaders import (
|
||||
# load_vposer,
|
||||
# load_state,
|
||||
# load_gmm,
|
||||
resolve_cfg_paths,
|
||||
)
|
||||
from util.logger import Logger
|
||||
from util.tensor import get_device, move_to, detach_all, to_torch
|
||||
|
||||
from run_vis import run_vis
|
||||
|
||||
import hydra
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
import time
|
||||
|
||||
N_STAGES = 3
|
||||
|
||||
import sys
|
||||
sys.path.append('src/human_body_prior')
|
||||
sys.path.append('HMP/')
|
||||
# print(sys.path)
|
||||
# print(torch.cuda.memory_summary())
|
||||
from HMP.fitting import run_prior
|
||||
|
||||
|
||||
def set_seed(seed=42):
|
||||
"""
|
||||
Set random seed for reproducibility
|
||||
"""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.backends.cudnn.deterministic = True
|
||||
os.environ['PYTHONHASHSEED'] = str(seed)
|
||||
|
||||
def run_opt(cfg, dataset, out_dir, device):
|
||||
a = time.time()
|
||||
args = cfg.data
|
||||
B = len(dataset)
|
||||
T = dataset.seq_len
|
||||
loader = DataLoader(dataset, batch_size=B, shuffle=False)
|
||||
|
||||
obs_data = move_to(next(iter(loader)), device)
|
||||
cam_data = move_to(dataset.get_camera_data(), device)
|
||||
print("Batch size (dataset_length), T (dataset.seq_len): ", B, len(dataset), T)
|
||||
print("OBS DATA", obs_data.keys())
|
||||
print("CAM DATA", cam_data.keys(), 'cam_R', cam_data['cam_R'])
|
||||
|
||||
# save cameras
|
||||
cam_R, cam_t = dataset.cam_data.cam2world()
|
||||
intrins = dataset.cam_data.intrins
|
||||
save_camera_json(f"cameras.json", cam_R, cam_t, intrins)
|
||||
|
||||
# check whether the cameras are static
|
||||
# if static, cannot optimize scale
|
||||
cfg.model.opt_scale &= not dataset.cam_data.is_static
|
||||
Logger.log(f"OPT SCALE {cfg.model.opt_scale}")
|
||||
|
||||
# loss weights for all stages
|
||||
all_loss_weights = cfg.optim.loss_weights
|
||||
assert all(len(wts) == N_STAGES for wts in all_loss_weights.values())
|
||||
stage_loss_weights = [
|
||||
{k: wts[i] for k, wts in all_loss_weights.items()} for i in range(N_STAGES)
|
||||
]
|
||||
max_loss_weights = {k: max(wts) for k, wts in all_loss_weights.items()}
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
cfg.paths.base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../"))
|
||||
paths = cfg.paths
|
||||
Logger.log(f"Loading pose prior from {paths.vposer}")
|
||||
# pose_prior, _ = load_model(paths.vposer, model_code=VPoser, remove_words_in_model_weights='vp_model.', disable_grad=True)
|
||||
# pose_prior = pose_prior.to(device)
|
||||
pose_prior = None
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
paths = cfg.paths
|
||||
Logger.log(f"Loading hand model")
|
||||
# Instantiate MANO model
|
||||
mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()}
|
||||
print('initializing MANO model with cfgs:', mano_cfg)
|
||||
hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device)
|
||||
|
||||
################################################################
|
||||
######################## optimization ##########################
|
||||
################################################################
|
||||
margs = cfg.model
|
||||
base_model = BaseSceneModel(
|
||||
B, T, hand_model, pose_prior, **margs
|
||||
)
|
||||
|
||||
base_model.initialize(obs_data, cam_data)
|
||||
base_model.to(device)
|
||||
|
||||
# save initial results for later visualization
|
||||
save_input_poses(dataset, os.path.join(out_dir, "hamer"), args.seq)
|
||||
save_initial_predictions(base_model, os.path.join(out_dir, "init"), args.seq)
|
||||
|
||||
opts = cfg.optim.options
|
||||
vis_scale = 0.25
|
||||
vis = None
|
||||
if opts.vis_every > 0:
|
||||
vis = init_viewer(
|
||||
dataset.img_size,
|
||||
cam_data["intrins"][0],
|
||||
vis_scale=vis_scale,
|
||||
bg_paths=dataset.sel_img_paths,
|
||||
fps=cfg.fps,
|
||||
)
|
||||
print("OPTIMIZER OPTIONS:", opts)
|
||||
|
||||
writer = SummaryWriter(out_dir)
|
||||
|
||||
print('start optimization')
|
||||
a = time.time()
|
||||
optim = RootOptimizer(base_model, stage_loss_weights, **opts)
|
||||
optim.run(obs_data, cfg.optim.root.num_iters, out_dir, vis, writer)
|
||||
|
||||
args = cfg.optim.smooth
|
||||
print(args)
|
||||
b = time.time()
|
||||
print('root optimization time: ', b - a)
|
||||
optim = SmoothOptimizer(
|
||||
base_model, stage_loss_weights, opt_scale=args.opt_scale, **opts
|
||||
)
|
||||
optim.run(obs_data, args.num_iters, out_dir, vis, writer)
|
||||
c = time.time()
|
||||
print('Smooth optimization time: ', c - b)
|
||||
|
||||
# HMP
|
||||
if cfg.run_prior and not glob.glob(os.path.join(out_dir, 'prior', '*_world_results.npz')):
|
||||
run_prior(cfg, dataset, out_dir, device, ['smooth_fit'], \
|
||||
obs_data, hand_model, cfg, cfg.data, os.path.join(out_dir, 'prior'))
|
||||
d = time.time()
|
||||
print('prior optimization time: ', d-c)
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="confs", config_name="config.yaml")
|
||||
def main(cfg: DictConfig):
|
||||
OmegaConf.register_new_resolver("eval", eval)
|
||||
print('run_opt.py: ', cfg)
|
||||
|
||||
# Set random seed
|
||||
set_seed(cfg.get('seed', 42))
|
||||
|
||||
out_dir = os.getcwd()
|
||||
print("out_dir", out_dir)
|
||||
Logger.init(f"{out_dir}/opt_log.txt")
|
||||
|
||||
# make sure we get all necessary inputs
|
||||
print("init SOURCES", cfg.data.sources)
|
||||
cfg.data.sources = expand_source_paths(cfg.data.sources)
|
||||
print("SOURCES", cfg.data.sources)
|
||||
|
||||
dataset = get_dataset_from_cfg(cfg)
|
||||
save_track_info(dataset, out_dir)
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = str(cfg.get("gpu"))
|
||||
print("CUDA_VISIBLE_DEVICES", os.environ["CUDA_VISIBLE_DEVICES"])
|
||||
device_id = cfg.get("gpu")
|
||||
|
||||
if cfg.run_opt:
|
||||
device = get_device(device_id)
|
||||
run_opt(cfg, dataset, out_dir, device)
|
||||
|
||||
if cfg.run_vis:
|
||||
run_vis(
|
||||
cfg, dataset, out_dir, device_id, **cfg.get("vis", dict())
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user