Spaces:
Build error
Build error
import torch | |
from torch.utils import data | |
import numpy as np | |
from os.path import join as pjoin | |
import random | |
import codecs as cs | |
from tqdm import tqdm | |
class VQMotionDataset(data.Dataset): | |
def __init__(self, dataset_name, feat_bias = 5, window_size = 64, unit_length = 8): | |
self.window_size = window_size | |
self.unit_length = unit_length | |
self.feat_bias = feat_bias | |
self.dataset_name = dataset_name | |
min_motion_len = 40 if dataset_name =='t2m' else 24 | |
if dataset_name == 't2m': | |
self.data_root = './dataset/HumanML3D' | |
self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') | |
self.text_dir = pjoin(self.data_root, 'texts') | |
self.joints_num = 22 | |
radius = 4 | |
fps = 20 | |
self.max_motion_length = 196 | |
dim_pose = 263 | |
self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta' | |
#kinematic_chain = paramUtil.t2m_kinematic_chain | |
elif dataset_name == 'kit': | |
self.data_root = './dataset/KIT-ML' | |
self.motion_dir = pjoin(self.data_root, 'new_joint_vecs') | |
self.text_dir = pjoin(self.data_root, 'texts') | |
self.joints_num = 21 | |
radius = 240 * 8 | |
fps = 12.5 | |
dim_pose = 251 | |
self.max_motion_length = 196 | |
self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta' | |
#kinematic_chain = paramUtil.kit_kinematic_chain | |
joints_num = self.joints_num | |
mean = np.load(pjoin(self.meta_dir, 'mean.npy')) | |
std = np.load(pjoin(self.meta_dir, 'std.npy')) | |
split_file = pjoin(self.data_root, 'train.txt') | |
data_dict = {} | |
id_list = [] | |
with cs.open(split_file, 'r') as f: | |
for line in f.readlines(): | |
id_list.append(line.strip()) | |
new_name_list = [] | |
length_list = [] | |
for name in tqdm(id_list): | |
try: | |
motion = np.load(pjoin(self.motion_dir, name + '.npy')) | |
if (len(motion)) < min_motion_len or (len(motion) >= 200): | |
continue | |
data_dict[name] = {'motion': motion, | |
'length': len(motion), | |
'name': name} | |
new_name_list.append(name) | |
length_list.append(len(motion)) | |
except: | |
# Some motion may not exist in KIT dataset | |
pass | |
self.mean = mean | |
self.std = std | |
self.length_arr = np.array(length_list) | |
self.data_dict = data_dict | |
self.name_list = new_name_list | |
def inv_transform(self, data): | |
return data * self.std + self.mean | |
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'] | |
m_length = (m_length // self.unit_length) * self.unit_length | |
idx = random.randint(0, len(motion) - m_length) | |
motion = motion[idx:idx+m_length] | |
"Z Normalization" | |
motion = (motion - self.mean) / self.std | |
return motion, name | |
def DATALoader(dataset_name, | |
batch_size = 1, | |
num_workers = 8, unit_length = 4) : | |
train_loader = torch.utils.data.DataLoader(VQMotionDataset(dataset_name, unit_length=unit_length), | |
batch_size, | |
shuffle=True, | |
num_workers=num_workers, | |
#collate_fn=collate_fn, | |
drop_last = True) | |
return train_loader | |
def cycle(iterable): | |
while True: | |
for x in iterable: | |
yield x | |