Spaces:
Running
Running
| import numpy as np | |
| import pickle | |
| from os.path import join as pjoin | |
| POS_enumerator = { | |
| 'VERB': 0, | |
| 'NOUN': 1, | |
| 'DET': 2, | |
| 'ADP': 3, | |
| 'NUM': 4, | |
| 'AUX': 5, | |
| 'PRON': 6, | |
| 'ADJ': 7, | |
| 'ADV': 8, | |
| 'Loc_VIP': 9, | |
| 'Body_VIP': 10, | |
| 'Obj_VIP': 11, | |
| 'Act_VIP': 12, | |
| 'Desc_VIP': 13, | |
| 'OTHER': 14, | |
| } | |
| Loc_list = ('left', 'right', 'clockwise', 'counterclockwise', 'anticlockwise', 'forward', 'back', 'backward', | |
| 'up', 'down', 'straight', 'curve') | |
| Body_list = ('arm', 'chin', 'foot', 'feet', 'face', 'hand', 'mouth', 'leg', 'waist', 'eye', 'knee', 'shoulder', 'thigh') | |
| Obj_List = ('stair', 'dumbbell', 'chair', 'window', 'floor', 'car', 'ball', 'handrail', 'baseball', 'basketball') | |
| Act_list = ('walk', 'run', 'swing', 'pick', 'bring', 'kick', 'put', 'squat', 'throw', 'hop', 'dance', 'jump', 'turn', | |
| 'stumble', 'dance', 'stop', 'sit', 'lift', 'lower', 'raise', 'wash', 'stand', 'kneel', 'stroll', | |
| 'rub', 'bend', 'balance', 'flap', 'jog', 'shuffle', 'lean', 'rotate', 'spin', 'spread', 'climb') | |
| Desc_list = ('slowly', 'carefully', 'fast', 'careful', 'slow', 'quickly', 'happy', 'angry', 'sad', 'happily', | |
| 'angrily', 'sadly') | |
| VIP_dict = { | |
| 'Loc_VIP': Loc_list, | |
| 'Body_VIP': Body_list, | |
| 'Obj_VIP': Obj_List, | |
| 'Act_VIP': Act_list, | |
| 'Desc_VIP': Desc_list, | |
| } | |
| class WordVectorizer(object): | |
| def __init__(self, meta_root, prefix): | |
| vectors = np.load(pjoin(meta_root, '%s_data.npy'%prefix)) | |
| words = pickle.load(open(pjoin(meta_root, '%s_words.pkl'%prefix), 'rb')) | |
| self.word2idx = pickle.load(open(pjoin(meta_root, '%s_idx.pkl'%prefix), 'rb')) | |
| self.word2vec = {w: vectors[self.word2idx[w]] for w in words} | |
| def _get_pos_ohot(self, pos): | |
| pos_vec = np.zeros(len(POS_enumerator)) | |
| if pos in POS_enumerator: | |
| pos_vec[POS_enumerator[pos]] = 1 | |
| else: | |
| pos_vec[POS_enumerator['OTHER']] = 1 | |
| return pos_vec | |
| def __len__(self): | |
| return len(self.word2vec) | |
| def __getitem__(self, item): | |
| word, pos = item.split('/') | |
| if word in self.word2vec: | |
| word_vec = self.word2vec[word] | |
| vip_pos = None | |
| for key, values in VIP_dict.items(): | |
| if word in values: | |
| vip_pos = key | |
| break | |
| if vip_pos is not None: | |
| pos_vec = self._get_pos_ohot(vip_pos) | |
| else: | |
| pos_vec = self._get_pos_ohot(pos) | |
| else: | |
| word_vec = self.word2vec['unk'] | |
| pos_vec = self._get_pos_ohot('OTHER') | |
| return word_vec, pos_vec | |
| class WordVectorizerV2(WordVectorizer): | |
| def __init__(self, meta_root, prefix): | |
| super(WordVectorizerV2, self).__init__(meta_root, prefix) | |
| self.idx2word = {self.word2idx[w]: w for w in self.word2idx} | |
| def __getitem__(self, item): | |
| word_vec, pose_vec = super(WordVectorizerV2, self).__getitem__(item) | |
| word, pos = item.split('/') | |
| if word in self.word2vec: | |
| return word_vec, pose_vec, self.word2idx[word] | |
| else: | |
| return word_vec, pose_vec, self.word2idx['unk'] | |
| def itos(self, idx): | |
| if idx == len(self.idx2word): | |
| return "pad" | |
| return self.idx2word[idx] |