157 lines
5.5 KiB
Python
157 lines
5.5 KiB
Python
import os
|
|
import rich
|
|
import random
|
|
import pickle
|
|
import codecs as cs
|
|
import numpy as np
|
|
from torch.utils import data
|
|
from rich.progress import track
|
|
from os.path import join as pjoin
|
|
|
|
|
|
class MotionDataset(data.Dataset):
|
|
def __init__(
|
|
self,
|
|
data_root,
|
|
split,
|
|
mean,
|
|
std,
|
|
max_motion_length=196,
|
|
min_motion_length=20,
|
|
unit_length=4,
|
|
fps=20,
|
|
tmpFile=True,
|
|
tiny=False,
|
|
debug=False,
|
|
**kwargs,
|
|
):
|
|
|
|
# restrian the length of motion and text
|
|
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())
|
|
|
|
# Debug mode
|
|
if tiny or debug:
|
|
enumerator = enumerate(
|
|
track(
|
|
self.id_list,
|
|
f"Loading HumanML3D {split}",
|
|
))
|
|
maxdata = 100
|
|
subset = '_tiny'
|
|
else:
|
|
enumerator = enumerate(self.id_list)
|
|
maxdata = 1e10
|
|
subset = ''
|
|
|
|
new_name_list = []
|
|
motion_dict = {}
|
|
|
|
# Fast loading
|
|
if os.path.exists(pjoin(data_root, f'tmp/{split}{subset}_motion.pkl')):
|
|
with rich.progress.open(pjoin(data_root, f'tmp/{split}{subset}_motion.pkl'),
|
|
'rb', description=f"Loading HumanML3D {split}") as file:
|
|
motion_dict = pickle.load(file)
|
|
with open(pjoin(data_root, f'tmp/{split}{subset}_index.pkl'), 'rb') as file:
|
|
new_name_list = pickle.load(file)
|
|
else:
|
|
for idx, name in enumerator:
|
|
if len(new_name_list) > maxdata:
|
|
break
|
|
try:
|
|
motion = [np.load(pjoin(motion_dir, name + ".npy"))]
|
|
|
|
# Read text
|
|
with cs.open(pjoin(text_dir, name + '.txt')) as f:
|
|
text_data = []
|
|
flag = False
|
|
lines = f.readlines()
|
|
|
|
for line in lines:
|
|
try:
|
|
line_split = line.strip().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
|
|
|
|
if f_tag == 0.0 and to_tag == 0.0:
|
|
flag = True
|
|
else:
|
|
motion_new = [tokens[int(f_tag*fps/unit_length) : int(to_tag*fps/unit_length)] for tokens in motion if int(f_tag*fps/unit_length) < int(to_tag*fps/unit_length)]
|
|
|
|
if len(motion_new) == 0:
|
|
continue
|
|
new_name = '%s_%f_%f'%(name, f_tag, to_tag)
|
|
|
|
motion_dict[new_name] = {
|
|
'motion': motion_new,
|
|
"length": [len(m[0]) for m in motion_new]}
|
|
new_name_list.append(new_name)
|
|
except:
|
|
pass
|
|
|
|
if flag:
|
|
motion_dict[name] = {
|
|
'motion': motion,
|
|
"length": [len(m[0]) for m in motion]}
|
|
new_name_list.append(name)
|
|
except:
|
|
pass
|
|
|
|
if tmpFile:
|
|
os.makedirs(pjoin(data_root, 'tmp'), exist_ok=True)
|
|
|
|
with open(pjoin(data_root, f'tmp/{split}{subset}_motion.pkl'),'wb') as file:
|
|
pickle.dump(motion_dict, file)
|
|
with open(pjoin(data_root, f'tmp/{split}{subset}_index.pkl'), 'wb') as file:
|
|
pickle.dump(new_name_list, file)
|
|
|
|
self.motion_dict = motion_dict
|
|
self.name_list = new_name_list
|
|
self.nfeats = motion_dict[new_name_list[0]]['motion'][0].shape[1]
|
|
|
|
def __len__(self):
|
|
return len(self.name_list)
|
|
|
|
def __getitem__(self, item):
|
|
data = self.motion_dict[self.name_list[item]]
|
|
motion_list, m_length = data["motion"], data["length"]
|
|
|
|
# Randomly select a motion
|
|
motion = random.choice(motion_list)
|
|
|
|
# Crop the motions in to times of 4, and introduce small variations
|
|
if self.unit_length < 10:
|
|
coin2 = np.random.choice(["single", "single", "double"])
|
|
else:
|
|
coin2 = "single"
|
|
|
|
if coin2 == "double":
|
|
m_length = (m_length // self.unit_length - 1) * self.unit_length
|
|
elif coin2 == "single":
|
|
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 None, motion, m_length, None, None, None, None,
|