Upload 8 files

dataset/dataset_TM_eval.py

@@ -0,0 +1,217 @@
+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
+
+
+def collate_fn(batch):
+    batch.sort(key=lambda x: x[3], reverse=True)
+    return default_collate(batch)
+
+
+'''For use of training text-2-motion generative model'''
+class Text2MotionDataset(data.Dataset):
+    def __init__(self, dataset_name, is_test, w_vectorizer, feat_bias = 5, max_text_len = 20, unit_length = 4):
+        
+        self.max_length = 20
+        self.pointer = 0
+        self.dataset_name = dataset_name
+        self.is_test = is_test
+        self.max_text_len = max_text_len
+        self.unit_length = unit_length
+        self.w_vectorizer = w_vectorizer
+        if dataset_name == 't2m':
+            self.data_root = './dataset/Sample1'
+            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
+            kinematic_chain = paramUtil.t2m_kinematic_chain
+            self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+        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
+            kinematic_chain = paramUtil.kit_kinematic_chain
+            self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+
+        mean = np.load(pjoin(self.meta_dir, 'mean.npy'))
+        std = np.load(pjoin(self.meta_dir, 'std.npy'))
+        
+        if is_test:
+            split_file = pjoin(self.data_root, 'test.txt')
+        else:
+            split_file = pjoin(self.data_root, 'val.txt')
+
+        min_motion_len = 40 if self.dataset_name =='t2m' else 24
+        # min_motion_len = 64
+
+        joints_num = self.joints_num
+
+        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
+                text_data = []
+                flag = False
+                with cs.open(pjoin(self.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*fps) : int(to_tag*fps)]
+                                if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200):
+                                    continue
+                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                while new_name in data_dict:
+                                    new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                data_dict[new_name] = {'motion': n_motion,
+                                                       'length': len(n_motion),
+                                                       'text':[text_dict]}
+                                new_name_list.append(new_name)
+                                length_list.append(len(n_motion))
+                            except:
+                                print(line_split)
+                                print(line_split[2], line_split[3], f_tag, to_tag, name)
+                                # break
+
+                if flag:
+                    data_dict[name] = {'motion': motion,
+                                       'length': len(motion),
+                                       'text': text_data}
+                    new_name_list.append(name)
+                    length_list.append(len(motion))
+            except Exception as e:
+                # print(e)
+                pass
+
+        name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1]))
+        self.mean = mean
+        self.std = std
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = name_list
+        self.reset_max_len(self.max_length)
+
+    def reset_max_len(self, length):
+        assert length <= self.max_motion_length
+        self.pointer = np.searchsorted(self.length_arr, length)
+        print("Pointer Pointing at %d"%self.pointer)
+        self.max_length = length
+
+    def inv_transform(self, data):
+        return data * self.std + self.mean
+
+    def forward_transform(self, data):
+        return (data - self.mean) / self.std
+
+    def __len__(self):
+        return len(self.data_dict) - self.pointer
+
+    def __getitem__(self, item):
+        idx = self.pointer + item
+        name = self.name_list[idx]
+        data = self.data_dict[name]
+        # data = self.data_dict[self.name_list[idx]]
+        motion, m_length, text_list = data['motion'], data['length'], data['text']
+        # Randomly select a caption
+        text_data = random.choice(text_list)
+        caption, tokens = text_data['caption'], text_data['tokens']
+
+        if len(tokens) < self.max_text_len:
+            # pad with "unk"
+            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+            sent_len = len(tokens)
+            tokens = tokens + ['unk/OTHER'] * (self.max_text_len + 2 - sent_len)
+        else:
+            # crop
+            tokens = tokens[:self.max_text_len]
+            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+            sent_len = len(tokens)
+        pos_one_hots = []
+        word_embeddings = []
+        for token in tokens:
+            word_emb, pos_oh = self.w_vectorizer[token]
+            pos_one_hots.append(pos_oh[None, :])
+            word_embeddings.append(word_emb[None, :])
+        pos_one_hots = np.concatenate(pos_one_hots, axis=0)
+        word_embeddings = np.concatenate(word_embeddings, axis=0)
+
+        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
+
+        if m_length < self.max_motion_length:
+            motion = np.concatenate([motion,
+                                     np.zeros((self.max_motion_length - m_length, motion.shape[1]))
+                                     ], axis=0)
+
+        return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, '_'.join(tokens), name
+
+
+
+
+def DATALoader(dataset_name, is_test,
+                batch_size, w_vectorizer,
+                num_workers = 8, unit_length = 4) : 
+    
+    val_loader = torch.utils.data.DataLoader(Text2MotionDataset(dataset_name, is_test, w_vectorizer, unit_length=unit_length),
+                                              batch_size,
+                                              shuffle = True,
+                                              num_workers=num_workers,
+                                              collate_fn=collate_fn,
+                                              drop_last = True)
+    return val_loader
+
+
+def cycle(iterable):
+    while True:
+        for x in iterable:
+            yield x

dataset/dataset_TM_train.py

@@ -0,0 +1,157 @@
+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
+
+
+def collate_fn(batch):
+    batch.sort(key=lambda x: x[3], reverse=True)
+    return default_collate(batch)
+
+
+'''For use of training text-2-motion generative model'''
+class Text2MotionDataset(data.Dataset):
+    def __init__(self, dataset_name, feat_bias=5, unit_length=4, codebook_size=1024, tokenizer_name=None):
+        self.max_length = 64
+        self.pointer = 0
+        self.dataset_name = dataset_name
+        self.unit_length = unit_length
+        self.mot_end_idx = codebook_size
+        self.mot_pad_idx = codebook_size + 1
+
+        print(f"Loading dataset: {dataset_name}")
+
+        if dataset_name == 't2m':
+            self.data_root = './dataset/Sample1'
+            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 = 26 if unit_length == 8 else 51
+            dim_pose = 263
+            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 = 26 if unit_length == 8 else 51
+            kinematic_chain = paramUtil.kit_kinematic_chain
+
+        split_file = pjoin(self.data_root, 'train.txt')
+
+        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(self.data_root, tokenizer_name, f'{name}.npy'))
+
+                with cs.open(pjoin(self.text_dir, f'{name}.txt')) as f:
+                    text_data = []
+                    flag = False
+                    lines = f.readlines()
+
+                    for line in lines:
+                        try:
+                            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:
+                                m_token_list_new = [tokens[int(f_tag * fps / unit_length): int(to_tag * fps / unit_length)] for tokens in m_token_list if int(f_tag * fps / unit_length) < int(to_tag * fps / unit_length)]
+
+                                if len(m_token_list_new) == 0:
+                                    continue
+                                new_name = 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)
+                        except:
+                            pass
+
+                if flag:
+                    data_dict[name] = {'m_token_list': m_token_list,
+                                       'text': text_data}
+                    new_name_list.append(name)
+            except:
+                pass
+
+        self.data_dict = data_dict
+        self.name_list = new_name_list
+
+        print(f"Dataset loaded. Number of samples: {len(self.data_dict)}")
+
+    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 = random.choice(m_token_list)
+
+        text_data = random.choice(text_list)
+        caption = text_data['caption']
+
+        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 + 1 < self.max_motion_length:
+            m_tokens = np.concatenate([m_tokens, np.ones((1), dtype=int) * self.mot_end_idx, np.ones((self.max_motion_length - 1 - m_tokens_len), dtype=int) * self.mot_pad_idx], axis=0)
+        else:
+            m_tokens = np.concatenate([m_tokens, np.ones((1), dtype=int) * self.mot_end_idx], axis=0)
+
+        return caption, m_tokens.reshape(-1), m_tokens_len
+
+
+def DATALoader(dataset_name,
+                batch_size, codebook_size, tokenizer_name, unit_length=4,
+                num_workers = 8) : 
+
+    train_loader = torch.utils.data.DataLoader(Text2MotionDataset(dataset_name, codebook_size = codebook_size, tokenizer_name = tokenizer_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
+
+

dataset/dataset_VQ.py

@@ -0,0 +1,109 @@
+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, window_size = 64, unit_length = 4):
+        self.window_size = window_size
+        self.unit_length = unit_length
+        self.dataset_name = dataset_name
+
+        if dataset_name == 't2m':
+            self.data_root = './dataset/Sample1'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 22
+            self.max_motion_length = 196
+            self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+
+        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
+
+            self.max_motion_length = 196
+            self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+        
+        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')
+
+        self.data = []
+        self.lengths = []
+        id_list = []
+        with cs.open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        for name in tqdm(id_list):
+            try:
+                motion = np.load(pjoin(self.motion_dir, name + '.npy'))
+                if motion.shape[0] < self.window_size:
+                    continue
+                self.lengths.append(motion.shape[0] - self.window_size)
+                self.data.append(motion)
+            except:
+                # Some motion may not exist in KIT dataset
+                pass
+
+            
+        self.mean = mean
+        self.std = std
+        print("Total number of motions {}".format(len(self.data)))
+
+    def inv_transform(self, data):
+        return data * self.std + self.mean
+    
+    def compute_sampling_prob(self) : 
+        
+        prob = np.array(self.lengths, dtype=np.float32)
+        prob /= np.sum(prob)
+        return prob
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, item):
+        motion = self.data[item]
+        
+        idx = random.randint(0, len(motion) - self.window_size)
+
+        motion = motion[idx:idx+self.window_size]
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        return motion
+
+def DATALoader(dataset_name,
+               batch_size,
+               num_workers = 8,
+               window_size = 64,
+               unit_length = 4):
+    
+    trainSet = VQMotionDataset(dataset_name, window_size=window_size, unit_length=unit_length)
+    prob = trainSet.compute_sampling_prob()
+    sampler = torch.utils.data.WeightedRandomSampler(prob, num_samples = len(trainSet) * 1000, replacement=True)
+    train_loader = torch.utils.data.DataLoader(trainSet,
+                                              batch_size,
+                                              shuffle=True,
+                                              #sampler=sampler,
+                                              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

dataset/dataset_tokenize.py

@@ -0,0 +1,117 @@
+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/Sample1'
+            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

dataset/prepare/download_extractor.sh

@@ -0,0 +1,15 @@
+rm -rf checkpoints
+mkdir checkpoints
+cd checkpoints
+echo -e "Downloading extractors"
+gdown --fuzzy https://drive.google.com/file/d/1o7RTDQcToJjTm9_mNWTyzvZvjTWpZfug/view
+gdown --fuzzy https://drive.google.com/file/d/1KNU8CsMAnxFrwopKBBkC8jEULGLPBHQp/view
+
+
+unzip t2m.zip
+unzip kit.zip
+
+echo -e "Cleaning\n"
+rm t2m.zip
+rm kit.zip
+echo -e "Downloading done!"

dataset/prepare/download_glove.sh

@@ -0,0 +1,9 @@
+echo -e "Downloading glove (in use by the evaluators)"
+gdown --fuzzy https://drive.google.com/file/d/1bCeS6Sh_mLVTebxIgiUHgdPrroW06mb6/view?usp=sharing
+rm -rf glove
+
+unzip glove.zip
+echo -e "Cleaning\n"
+rm glove.zip
+
+echo -e "Downloading done!"

dataset/prepare/download_model.sh

@@ -0,0 +1,12 @@
+
+mkdir -p pretrained
+cd pretrained/
+
+echo -e "The pretrained model files will be stored in the 'pretrained' folder\n"
+gdown 1LaOvwypF-jM2Axnq5dc-Iuvv3w_G-WDE
+
+unzip VQTrans_pretrained.zip
+echo -e "Cleaning\n"
+rm VQTrans_pretrained.zip
+
+echo -e "Downloading done!"

dataset/prepare/download_smpl.sh

@@ -0,0 +1,13 @@
+
+mkdir -p body_models
+cd body_models/
+
+echo -e "The smpl files will be stored in the 'body_models/smpl/' folder\n"
+gdown 1INYlGA76ak_cKGzvpOV2Pe6RkYTlXTW2
+rm -rf smpl
+
+unzip smpl.zip
+echo -e "Cleaning\n"
+rm smpl.zip
+
+echo -e "Downloading done!"