import torch from torch.utils.data import Dataset from tqdm import tqdm from utils import dist_util import numpy as np from data_loaders.tensors import lengths_to_mask import os import json from data_loaders.amass.babel import get_tokens from data_loaders.humanml.utils.get_opt import get_opt from data_loaders.humanml.data.dataset import process_tokens from os.path import join as pjoin def pad_sample_with_zeros(sample, max_len=250): # pad inp, change lenghts, and pad is transition seq_len, n_feats = sample.shape len_to_pad = max_len - seq_len np.zeros_like(sample) sample_padding = np.zeros((len_to_pad, n_feats)) sample = np.concatenate((sample, sample_padding)) return sample class CompMDMGeneratedDataset(Dataset): def load_model_kwargs_dataset(self, eval_file, scenario=""): import json with open(eval_file, 'r') as f: all_model_kwargs = json.load(f) print(f"loaded {eval_file}", len(all_model_kwargs)) # convert all "lengths" to torch final_model_kwargs = [] for i in range(len(all_model_kwargs)): idx = int(all_model_kwargs[i]['id']) kwargs = {"y": all_model_kwargs[i]} if scenario != "" and scenario is not None and "scenario" in kwargs['y'] and kwargs['y']['scenario'] != scenario: continue # skip this one kwargs['y']['lengths'] = torch.tensor([int(v) for v in kwargs['y']['lengths']]) kwargs['y']['mask'] = lengths_to_mask(kwargs['y']['lengths'], kwargs['y']['lengths'].max()).unsqueeze(1).unsqueeze(1) final_model_kwargs.append((idx, kwargs)) assert len(final_model_kwargs) > 0, f"No model kwargs found for this scenario: {scenario}" print(f"loaded {len(final_model_kwargs)} model kwargs for scenario {scenario if scenario != '' else '> all <'}") return iter(final_model_kwargs) def process_tokens(self, tokens): return process_tokens(tokens, self.opt.max_text_len, self.w_vectorizer) def __init__(self, args, model, diffusion, mm_num_samples, mm_num_repeats, eval_file): dataloader = self.load_model_kwargs_dataset(eval_file) #assert mm_num_samples < len(dataset) use_ddim = False # FIXME - hardcoded clip_denoised = False # FIXME - hardcoded sample_fn = ( diffusion.p_sample_loop if not use_ddim else diffusion.ddim_sample_loop ) generated_motion = [] mm_generated_motions = [] num_seqs = 32 # ALWAYS if mm_num_samples > 0: mm_idxs = np.random.choice(len(dataloader), mm_num_samples // num_seqs +1, replace=False) mm_idxs = np.sort(mm_idxs) else: mm_idxs = [] print('mm_idxs', mm_idxs) model.eval() with torch.no_grad(): for i, (motion, model_kwargs) in tqdm(enumerate(dataloader)): is_mm = i in mm_idxs repeat_times = mm_num_repeats if is_mm else 1 mm_motions = [] for t in range(repeat_times): sample = sample_fn( motion.shape, clip_denoised=clip_denoised, model_kwargs=model_kwargs, #skip_timesteps=0, # 0 is the default value - i.e. don't skip any step progress=False, ) if t == 0: sub_dicts = [{'motion': sample[bs_i].squeeze().permute(1,0).cpu().numpy(), 'length': model_kwargs['y']['lengths'][bs_i].cpu().numpy(), 'caption': model_kwargs['y']['text'][bs_i], } for bs_i in range(num_seqs)] generated_motion += sub_dicts if is_mm: mm_motions += [{'motion': sample[bs_i].squeeze().permute(1, 0).cpu().numpy(), 'length': model_kwargs['y']['lengths'][bs_i].cpu().numpy(), } for bs_i in range(num_seqs)] if is_mm: mm_generated_motions += [{ 'caption': model_kwargs['y']['text'][bs_i], 'mm_motions': mm_motions[bs_i::num_seqs], # collect all 10 repeats from the (32*10) generated motions } for bs_i in range(num_seqs)] self.generated_motion = generated_motion self.mm_generated_motion = mm_generated_motions def __len__(self): return len(self.generated_motion) def __getitem__(self, item): data = self.generated_motion[item] motion, m_length, caption = data['motion'], data['length'], data['caption'] if self.dataset_name != "babel": # babel already takes care of its de)normalization itself normed_motion = motion denormed_motion = normed_motion * self.std + self.mean renormed_motion = (denormed_motion - self.mean_for_eval) / self.std_for_eval # according to T2M norms motion = renormed_motion # This step is needed because T2M evaluators expect their norm convention tokens = get_tokens(caption) word_embeddings, pos_one_hots, sent_len, tokens = self.process_tokens(tokens) return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, tokens def inv_transform(self, data): return data * self.std_for_eval + self.mean_for_eval def switch_target(self, target): assert target in ['motion', ], "Only motion eval target is available for non-unfolding dataset" class CompMDMUnfoldingGeneratedDataset(CompMDMGeneratedDataset): def __init__(self, args, model, diffusion, max_motion_length, eval_file, w_vectorizer=None, opt=None, precomputed_folder=None, scenario=""): self.dataset_name = args.dataset self.w_vectorizer = w_vectorizer self.opt = opt assert self.dataset_name == "babel" or self.dataset_name == "humanml", "Only babel and humanml are supported" if self.dataset_name == "humanml": self.mean = np.load(pjoin(opt.data_root, 'Mean.npy')) self.std = np.load(pjoin(opt.data_root, 'Std.npy')) self.mean_for_eval = np.load(pjoin(opt.meta_dir, f'{opt.dataset_name}_mean.npy')) self.std_for_eval = np.load(pjoin(opt.meta_dir, f'{opt.dataset_name}_std.npy')) dataloader = self.load_model_kwargs_dataset(eval_file, scenario=scenario) self.max_motion_length = max_motion_length # Will be changed later by the evaluation script for each copy of this dataset self.step_to_eval = 1 self.transition = False generated_transitions = [] generated_motion = [] os.makedirs(precomputed_folder, exist_ok=True) with torch.no_grad(): for i, model_kwargs in tqdm(dataloader): precomputed_file_kwargs = os.path.join(precomputed_folder, f'{i:02d}_kwargs.json') precomputed_file_pt = os.path.join(precomputed_folder, f'{i:02d}.pt') precomputed_file_npy = os.path.join(precomputed_folder, f'{i:02d}.npy') if os.path.exists(precomputed_file_kwargs): # load it from precomputed file #print(f'Loading precomputed file {precomputed_file}') loaded_kwargs = json.load(open(precomputed_file_kwargs, 'r')) kwargs = loaded_kwargs if "y" in loaded_kwargs else {"y": loaded_kwargs} # to keep compatibility with the old format # assert equal assert kwargs['y']['lengths'] == model_kwargs['y']['lengths'].tolist() assert kwargs['y']['text'] == model_kwargs['y']['text'] if os.path.exists(precomputed_file_pt): unfolded = torch.load(precomputed_file_pt) elif os.path.exists(precomputed_file_npy): # if unfolded contains rots and transl, we need to convert them to our features representation unfolded = np.load(precomputed_file_npy, allow_pickle=True).item() rots = torch.tensor(unfolded["rots"]) # --> [seq_len, 22, 3, 3], rot matrices transl = torch.tensor(unfolded["transl"]) # --> [seq_len, 3], translations # we need to go from 3x3 matrices to axis angle, and then we got it :) from utils.rotation_conversions import matrix_to_axis_angle rots = matrix_to_axis_angle(rots) # --> [seq_len, 22, 3], axis angle from data_loaders.amass.tools.smpl import smpl_data_to_matrix_and_trans smpl_data = { "poses": rots, "trans": transl, } smpl_data = smpl_data_to_matrix_and_trans(smpl_data, nohands=True) from data_loaders.amass.transforms import SlimSMPLTransform transform = SlimSMPLTransform(batch_size=32, name='SlimSMPLTransform', ename='smplnh', normalization=True) features = transform.rots2rfeats(smpl_data) unfolded = features.permute(1, 0).unsqueeze(1).unsqueeze(0) torch.save(unfolded, precomputed_file_pt) else: assert False, "Precomputed file not found" else: assert model is not None and diffusion is not None, "Model and diffusion must be provided for evaluation if precomputed files are not available" # compute it model.eval() unfolded = diffusion.p_sample_loop( clip_denoised=False, model_kwargs=model_kwargs, progress=False#True, ) # --> [1, 263, 1, BS * single_lengths], where single_lengths is each length in model_kwargs['y']['lengths'] # store it #print(f'Storing precomputed file {precomputed_file}') torch.save(unfolded, precomputed_file_pt) # save in float32 json.dump({'lengths': model_kwargs['y']['lengths'].tolist(), 'text': model_kwargs['y']['text']}, open(precomputed_file_kwargs, 'w')) # evaluates on the whole motion segment start = 0 num_seqs = len(model_kwargs['y']['lengths']) for bs_i in range(num_seqs): end = (start + model_kwargs['y']['lengths'][bs_i]) if bs_i != num_seqs - 1 else None motion_slice = unfolded[..., start:end].squeeze().permute(1, 0).cpu().numpy() assert motion_slice.shape[0] == model_kwargs['y']['lengths'][bs_i], f'{motion_slice.shape[0]} != {model_kwargs["y"]["lengths"][bs_i]}' generated_motion.append({ 'motion': pad_sample_with_zeros(motion_slice, self.max_motion_length), 'length': model_kwargs['y']['lengths'][bs_i], 'caption': model_kwargs['y']['text'][bs_i], }) start = end # only keeps the transition l_margin = (args.transition_length // 2) r_margin = args.transition_length - l_margin mid = 0 for bs_i in range(num_seqs - 1): # last one has no transition mid += model_kwargs['y']['lengths'][bs_i] motion_slice = unfolded[..., mid - l_margin : mid + r_margin].squeeze().permute(1, 0).cpu().numpy() assert motion_slice.shape[0] == args.transition_length generated_transitions.append({ 'motion': motion_slice, # all transitions with length args.transition_length, so no need for padding 'length': args.transition_length, 'caption': model_kwargs['y']['text'][bs_i], }) self.generated_inbetween = generated_motion self.generated_transitions = generated_transitions self.switch_target('motion') # 'motion' or 'transition' def switch_target(self, target): """ Switches between 'motion' and 'transition' targets. In 'motion' target, the dataset returns the full motion segment. In 'transition' target, the dataset returns only the transition part of the motion """ assert target in ['motion', 'transition'] self.target = target self.generated_motion = self.generated_inbetween if self.target == 'motion' else self.generated_transitions