| 14 | |
| 15 | |
| 16 | class EvaluationDataset(Dataset): |
| 17 | |
| 18 | def __init__(self, opt, trainer, dataset, w_vectorizer, mm_num_samples, mm_num_repeats): |
| 19 | assert mm_num_samples < len(dataset) |
| 20 | print(opt.model_dir) |
| 21 | |
| 22 | dataloader = DataLoader(dataset, batch_size=1, num_workers=1, shuffle=True) |
| 23 | epoch, it = trainer.load(pjoin(opt.model_dir, opt.which_epoch + '.tar')) |
| 24 | |
| 25 | generated_motion = [] |
| 26 | min_mov_length = 10 if opt.dataset_name == 't2m' else 6 |
| 27 | |
| 28 | trainer.eval_mode() |
| 29 | trainer.to(opt.device) |
| 30 | |
| 31 | # Pre-process all target captions |
| 32 | mm_generated_motions = [] |
| 33 | mm_idxs = np.random.choice(len(dataset), mm_num_samples, replace=False) |
| 34 | mm_idxs = np.sort(mm_idxs) |
| 35 | all_caption = [] |
| 36 | all_m_lens = [] |
| 37 | all_data = [] |
| 38 | with torch.no_grad(): |
| 39 | for i, data in tqdm(enumerate(dataloader)): |
| 40 | word_emb, pos_ohot, caption, cap_lens, motions, m_lens, tokens = data |
| 41 | all_data.append(data) |
| 42 | tokens = tokens[0].split('_') |
| 43 | mm_num_now = len(mm_generated_motions) |
| 44 | is_mm = True if ((mm_num_now < mm_num_samples) and (i == mm_idxs[mm_num_now])) else False |
| 45 | repeat_times = mm_num_repeats if is_mm else 1 |
| 46 | m_lens = max(m_lens // opt.unit_length * opt.unit_length, min_mov_length * opt.unit_length) |
| 47 | m_lens = min(m_lens, opt.max_motion_length) |
| 48 | if isinstance(m_lens, int): |
| 49 | m_lens = torch.LongTensor([m_lens]).to(opt.device) |
| 50 | else: |
| 51 | m_lens = m_lens.to(opt.device) |
| 52 | for t in range(repeat_times): |
| 53 | all_m_lens.append(m_lens) |
| 54 | all_caption.extend(caption) |
| 55 | if is_mm: |
| 56 | mm_generated_motions.append(0) |
| 57 | all_m_lens = torch.stack(all_m_lens) |
| 58 | |
| 59 | # Generate all sequences |
| 60 | with torch.no_grad(): |
| 61 | all_pred_motions = trainer.generate(all_caption, all_m_lens, opt.dim_pose) |
| 62 | |
| 63 | cur_idx = 0 |
| 64 | mm_generated_motions = [] |
| 65 | with torch.no_grad(): |
| 66 | for i, data_dummy in tqdm(enumerate(dataloader)): |
| 67 | data = all_data[i] |
| 68 | word_emb, pos_ohot, caption, cap_lens, motions, m_lens, tokens = data |
| 69 | tokens = tokens[0].split('_') |
| 70 | mm_num_now = len(mm_generated_motions) |
| 71 | is_mm = True if ((mm_num_now < mm_num_samples) and (i == mm_idxs[mm_num_now])) else False |
| 72 | repeat_times = mm_num_repeats if is_mm else 1 |
| 73 | mm_motions = [] |