| import random |
| import numpy as np |
| from torch.utils import data |
| from .dataset_t2m import Text2MotionDataset |
| import codecs as cs |
| from os.path import join as pjoin |
|
|
|
|
| class Text2MotionDatasetM2T(data.Dataset): |
|
|
| def __init__( |
| self, |
| data_root, |
| split, |
| mean, |
| std, |
| max_motion_length=196, |
| min_motion_length=40, |
| unit_length=4, |
| fps=20, |
| tmpFile=True, |
| tiny=False, |
| debug=False, |
| **kwargs, |
| ): |
| |
| self.max_motion_length = max_motion_length |
| self.min_motion_length = min_motion_length |
| self.unit_length = unit_length |
| |
| |
| self.mean = mean |
| self.std = std |
| |
| |
| split_file = pjoin(data_root, split + '.txt') |
| motion_dir = pjoin(data_root, 'new_joint_vecs') |
| text_dir = pjoin(data_root, 'texts') |
|
|
| |
| self.id_list = [] |
| with cs.open(split_file, "r") as f: |
| for line in f.readlines(): |
| self.id_list.append(line.strip()) |
| |
| new_name_list = [] |
| length_list = [] |
| data_dict = {} |
| for name in self.id_list: |
| |
| motion = np.load(pjoin(motion_dir, name + '.npy')) |
| if (len(motion)) < self.min_motion_length or (len(motion) >= 200): |
| continue |
| |
| |
| text_data = [] |
| flag = False |
| |
| with cs.open(pjoin(text_dir, name + '.txt')) as f: |
| for line in f.readlines(): |
| text_dict = {} |
| line_split = line.strip().split('#') |
| caption = line_split[0] |
| tokens = line_split[1].split(' ') |
| f_tag = float(line_split[2]) |
| to_tag = float(line_split[3]) |
| f_tag = 0.0 if np.isnan(f_tag) else f_tag |
| to_tag = 0.0 if np.isnan(to_tag) else to_tag |
|
|
| text_dict['caption'] = caption |
| text_dict['tokens'] = tokens |
| if f_tag == 0.0 and to_tag == 0.0: |
| flag = True |
| text_data.append(text_dict) |
| else: |
| try: |
| n_motion = motion[int(f_tag*20) : int(to_tag*20)] |
|
|
| if (len(n_motion)) < min_motion_length or (len(n_motion) >= 200): |
| continue |
| |
| new_name = "%s_%f_%f"%(name, f_tag, to_tag) |
| data_dict[new_name] = {'motion': n_motion, |
| 'length': len(n_motion), |
| 'text':[text_dict]} |
| new_name_list.append(new_name) |
| except: |
| print(line_split) |
| print(line_split[2], line_split[3], f_tag, to_tag, name) |
| if flag: |
| data_dict[name] = {'motion': motion, |
| 'length': len(motion), |
| 'name': name, |
| 'text': text_data} |
| |
| new_name_list.append(name) |
| length_list.append(len(motion)) |
| |
| |
| |
|
|
| self.length_arr = np.array(length_list) |
| self.data_dict = data_dict |
| self.name_list = new_name_list |
| self.nfeats = motion.shape[-1] |
| |
| |
| def __len__(self): |
| return len(self.data_dict) |
| |
| def __getitem__(self, item): |
| name = self.name_list[item] |
| data = self.data_dict[name] |
| motion, m_length = data['motion'], data['length'] |
|
|
| "Z Normalization" |
| motion = (motion - self.mean) / self.std |
|
|
| return name, motion, m_length, True, True, True, True, True, True |
|
|
