|
|
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 |
|
|
import utils.paramUtil as paramUtil |
|
|
from torch.utils.data._utils.collate import default_collate |
|
|
import os |
|
|
|
|
|
def collate_fn(batch): |
|
|
batch.sort(key=lambda x: x[2], reverse=True) |
|
|
return default_collate(batch) |
|
|
|
|
|
|
|
|
'''For use of training text-2-motion generative model''' |
|
|
class Text2MotionDataset(data.Dataset): |
|
|
def __init__(self, dataset_name, unit_length = 4, latent_dir=None): |
|
|
|
|
|
self.max_length = 64 |
|
|
self.pointer = 0 |
|
|
self.dataset_name = dataset_name |
|
|
self.unit_length = unit_length |
|
|
|
|
|
if dataset_name == 't2m_272': |
|
|
self.data_root = './humanml3d_272' |
|
|
self.text_dir = pjoin(self.data_root, 'texts') |
|
|
self.joints_num = 22 |
|
|
fps = 30 |
|
|
self.max_motion_length = 78 |
|
|
dim_pose = 272 |
|
|
split_file = pjoin(self.data_root, 'split', 'train.txt') |
|
|
|
|
|
else: |
|
|
raise ValueError(f"Dataset {dataset_name} not supported") |
|
|
|
|
|
id_list = [] |
|
|
with cs.open(split_file, 'r') as f: |
|
|
for line in f.readlines(): |
|
|
id_list.append(line.strip()) |
|
|
|
|
|
new_name_list = [] |
|
|
data_dict = {} |
|
|
for name in tqdm(id_list): |
|
|
try: |
|
|
m_token_list = np.load(pjoin(latent_dir, '%s.npy'%name)) |
|
|
except: |
|
|
continue |
|
|
|
|
|
|
|
|
with cs.open(pjoin(self.text_dir, name + '.txt')) as f: |
|
|
text_data = [] |
|
|
flag = False |
|
|
lines = f.readlines() |
|
|
|
|
|
for line in lines: |
|
|
text_dict = {} |
|
|
line_split = line.strip().split('#') |
|
|
caption = line_split[0] |
|
|
t_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'] = t_tokens |
|
|
|
|
|
if f_tag == 0.0 and to_tag == 0.0: |
|
|
flag = True |
|
|
text_data.append(text_dict) |
|
|
else: |
|
|
if int(f_tag*fps/unit_length) < int(to_tag*fps/unit_length): |
|
|
m_token_list_new = [m_token_list[int(f_tag*fps/unit_length) : int(to_tag*fps/unit_length)]] |
|
|
|
|
|
if len(m_token_list_new) == 0: |
|
|
continue |
|
|
|
|
|
new_name = '%s_%f_%f'%(name, f_tag, to_tag) |
|
|
|
|
|
data_dict[new_name] = {'m_token_list': m_token_list_new, |
|
|
'text':[text_dict]} |
|
|
new_name_list.append(new_name) |
|
|
|
|
|
if flag: |
|
|
data_dict[name] = {'m_token_list': m_token_list, |
|
|
'text':text_data} |
|
|
new_name_list.append(name) |
|
|
|
|
|
self.data_dict = data_dict |
|
|
self.name_list = new_name_list |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.data_dict) |
|
|
|
|
|
def __getitem__(self, item): |
|
|
data = self.data_dict[self.name_list[item]] |
|
|
m_token_list, text_list = data['m_token_list'], data['text'] |
|
|
m_tokens = np.array(m_token_list) |
|
|
|
|
|
text_data = random.choice(text_list) |
|
|
caption= text_data['caption'] |
|
|
|
|
|
if len(m_tokens.shape) == 3: |
|
|
m_tokens = m_tokens.squeeze(0) |
|
|
coin = np.random.choice([False, False, True]) |
|
|
if coin: |
|
|
coin2 = np.random.choice([True, False]) |
|
|
if coin2: |
|
|
m_tokens = m_tokens[:-1] |
|
|
else: |
|
|
m_tokens = m_tokens[1:] |
|
|
m_tokens_len = m_tokens.shape[0] |
|
|
|
|
|
if m_tokens_len < self.max_motion_length: |
|
|
m_tokens = np.concatenate([m_tokens, np.zeros((self.max_motion_length-m_tokens_len, m_tokens.shape[1]), dtype=int)], axis=0) |
|
|
return caption, m_tokens, m_tokens_len |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def DATALoader(dataset_name, |
|
|
batch_size, latent_dir, unit_length=4, |
|
|
num_workers = 8) : |
|
|
|
|
|
train_loader = torch.utils.data.DataLoader(Text2MotionDataset(dataset_name, latent_dir = latent_dir, unit_length=unit_length), |
|
|
batch_size, |
|
|
shuffle=True, |
|
|
num_workers=num_workers, |
|
|
drop_last = True) |
|
|
|
|
|
return train_loader |
|
|
|
|
|
|
|
|
def cycle(iterable): |
|
|
while True: |
|
|
for x in iterable: |
|
|
yield x |
|
|
|
|
|
|
