Spaces:
Runtime error
Runtime error
File size: 4,252 Bytes
4409449 |
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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
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
# Data mean and std
self.mean = mean
self.std = std
# Data path
split_file = pjoin(data_root, split + '.txt')
motion_dir = pjoin(data_root, 'new_joint_vecs')
text_dir = pjoin(data_root, 'texts')
# Data id list
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:
# try:
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))
# except:
# # Some motion may not exist in KIT dataset
# pass
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
|