mingyuan/MotionDiffuse

1import numpy as np2import pickle3from os.path import join as pjoin4 5POS_enumerator = {6    'VERB': 0,7    'NOUN': 1,8    'DET': 2,9    'ADP': 3,10    'NUM': 4,11    'AUX': 5,12    'PRON': 6,13    'ADJ': 7,14    'ADV': 8,15    'Loc_VIP': 9,16    'Body_VIP': 10,17    'Obj_VIP': 11,18    'Act_VIP': 12,19    'Desc_VIP': 13,20    'OTHER': 14,21}22 23Loc_list = ('left', 'right', 'clockwise', 'counterclockwise', 'anticlockwise', 'forward', 'back', 'backward',24            'up', 'down', 'straight', 'curve')25 26Body_list = ('arm', 'chin', 'foot', 'feet', 'face', 'hand', 'mouth', 'leg', 'waist', 'eye', 'knee', 'shoulder', 'thigh')27 28Obj_List = ('stair', 'dumbbell', 'chair', 'window', 'floor', 'car', 'ball', 'handrail', 'baseball', 'basketball')29 30Act_list = ('walk', 'run', 'swing', 'pick', 'bring', 'kick', 'put', 'squat', 'throw', 'hop', 'dance', 'jump', 'turn',31            'stumble', 'dance', 'stop', 'sit', 'lift', 'lower', 'raise', 'wash', 'stand', 'kneel', 'stroll',32            'rub', 'bend', 'balance', 'flap', 'jog', 'shuffle', 'lean', 'rotate', 'spin', 'spread', 'climb')33 34Desc_list = ('slowly', 'carefully', 'fast', 'careful', 'slow', 'quickly', 'happy', 'angry', 'sad', 'happily',35             'angrily', 'sadly')36 37VIP_dict = {38    'Loc_VIP': Loc_list,39    'Body_VIP': Body_list,40    'Obj_VIP': Obj_List,41    'Act_VIP': Act_list,42    'Desc_VIP': Desc_list,43}44 45 46class WordVectorizer(object):47    def __init__(self, meta_root, prefix):48        vectors = np.load(pjoin(meta_root, '%s_data.npy'%prefix))49        words = pickle.load(open(pjoin(meta_root, '%s_words.pkl'%prefix), 'rb'))50        word2idx = pickle.load(open(pjoin(meta_root, '%s_idx.pkl'%prefix), 'rb'))51        self.word2vec = {w: vectors[word2idx[w]] for w in words}52 53    def _get_pos_ohot(self, pos):54        pos_vec = np.zeros(len(POS_enumerator))55        if pos in POS_enumerator:56            pos_vec[POS_enumerator[pos]] = 157        else:58            pos_vec[POS_enumerator['OTHER']] = 159        return pos_vec60 61    def __len__(self):62        return len(self.word2vec)63 64    def __getitem__(self, item):65        word, pos = item.split('/')66        if word in self.word2vec:67            word_vec = self.word2vec[word]68            vip_pos = None69            for key, values in VIP_dict.items():70                if word in values:71                    vip_pos = key72                    break73            if vip_pos is not None:74                pos_vec = self._get_pos_ohot(vip_pos)75            else:76                pos_vec = self._get_pos_ohot(pos)77        else:78            word_vec = self.word2vec['unk']79            pos_vec = self._get_pos_ohot('OTHER')80        return word_vec, pos_vec