File size: 3,605 Bytes
8b3f571 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import sys, os
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), '..'))
import torch
from mmengine.registry import MODELS, METRICS
from mmengine.runner import Runner
from mmengine.config import Config
import torch.nn.functional as F
from src.utils.motion import recover_from_local_position_fast, recover_from_local_rotation
from tqdm import tqdm
import numpy as np
def vis_motion(motion, mean, std):
from src.metrics.recons_metric import recover_from_local_position
from src.utils.visualize import vis_3d_motion
motion = (motion * std + mean).numpy()
positions = recover_from_local_position(motion, 22) #
vis_3d_motion(positions[None], None, [f'video.mp4'], fps=30)
cfg = Config.fromfile('configs/eval_text.py')
ckp = torch.load('work_dirs/anymo_stage_3_3B_210e_48gpu/e210.pth')['state_dict']
lr = 2e-4
opt_iter = 400
# build model
model = MODELS.build(cfg.model)
model.load_state_dict(ckp)
model.eval().cuda()
# build dataloader
cfg.test_dataloader.batch_size = 2
dataloader = Runner.build_dataloader(cfg.test_dataloader)
# build evaluator
evaluator = METRICS.build(cfg.test_evaluator)
for data in tqdm(dataloader):
# prepare data
# import ipdb; ipdb.set_trace()
data = {k: v.cuda() if isinstance(v, torch.Tensor) else v for k, v in data.items()}
m_lens = data['motion_length']
caps = data['text_data']
_, _, ids = model(**data, mode='predict', return_ids=True)
# import ipdb; ipdb.set_trace()
ids_mask = (ids > -1).unsqueeze(1)
ids[ids < 0] = 0
emb = model.vqvae.quantizer.dequantize(ids.permute(1, 2, 0)) # B, T, C
emb = emb.detach().clone().permute(0, 2, 1) # B, C, T
emb = emb * ids_mask[0][0].unsqueeze(1)
emb.requires_grad = True
optimizer = torch.optim.AdamW([emb], lr=lr, betas=(0.5, 0.9), weight_decay=1e-6)
m_mask = torch.arange(max(m_lens)).cuda().expand(len(m_lens), max(m_lens)) < m_lens.unsqueeze(1)
pbar = tqdm(total=opt_iter)
for i in range(opt_iter):
pred_motion = model.vqvae.decoder(emb, ids_mask[0].squeeze(0).sum(1)).permute(0, 2, 1)
pred_traj = recover_from_local_position_fast(pred_motion * data['std'] + data['mean'], njoint=22)[:, :, 0]
gt_traj = recover_from_local_position_fast(data['motion'] * data['std'] + data['mean'], njoint=22)[:, :, 0]
tta_loss = F.mse_loss(pred_traj, gt_traj, reduction='none')
tta_loss = (tta_loss * m_mask.unsqueeze(-1)).sum() / (m_mask.sum() * tta_loss.shape[-1])
optimizer.zero_grad()
tta_loss.backward()
optimizer.step()
pbar.update(1)
pbar.set_description(f'tta_loss: {tta_loss.detach().cpu().item():.5f}')
pred_motions = model.vqvae.decoder(emb, ids_mask[0].squeeze(0).sum(1)).permute(0, 2, 1)
pred_motions = [pred_motions[i][:m_lens[i]] for i in range(pred_motions.shape[0])]
import ipdb; ipdb.set_trace()
# idx = 25;pred_data = recover_from_local_rotation((pred_motions[idx].cpu() * data['std'].cpu() + data['mean'].cpu()).detach().cpu().numpy()[:m_lens[idx]], njoint=22); np.savez('tmp.npz', **pred_data)
# idx = 25;pred_data = recover_from_local_rotation((data['motion'].cpu() * data['std'].cpu() + data['mean'].cpu()).cpu().numpy()[idx][:m_lens[idx]], njoint=22); np.savez('tmp.npz', **pred_data)
# idx = 25;vis_motion(pred_motions[idx].detach().cpu(), data['mean'].cpu(), data['std'].cpu())
# idx = 25;vis_motion(data['motion'][idx][:m_lens[idx]].cpu(), data['mean'].cpu(), data['std'].cpu())
# evaluator.process(data, (pred_motions, m_lens))
# print(evaluator.compute_metrics(evaluator.results)) |