radames/Text2Human-API
1
1from curses import A_ATTRIBUTES2 3import numpy4import torch5from pip import main6from sentence_transformers import SentenceTransformer, util7 8# predefined shape text9upper_length_text = [10 'sleeveless', 'without sleeves', 'sleeves have been cut off', 'tank top',11 'tank shirt', 'muscle shirt', 'short-sleeve', 'short sleeves',12 'with short sleeves', 'medium-sleeve', 'medium sleeves',13 'with medium sleeves', 'sleeves reach elbow', 'long-sleeve',14 'long sleeves', 'with long sleeves'15]16upper_length_attr = {17 'sleeveless': 0,18 'without sleeves': 0,19 'sleeves have been cut off': 0,20 'tank top': 0,21 'tank shirt': 0,22 'muscle shirt': 0,23 'short-sleeve': 1,24 'with short sleeves': 1,25 'short sleeves': 1,26 'medium-sleeve': 2,27 'with medium sleeves': 2,28 'medium sleeves': 2,29 'sleeves reach elbow': 2,30 'long-sleeve': 3,31 'long sleeves': 3,32 'with long sleeves': 333}34lower_length_text = [35 'three-point', 'medium', 'short', 'covering knee', 'cropped',36 'three-quarter', 'long', 'slack', 'of long length'37]38lower_length_attr = {39 'three-point': 0,40 'medium': 1,41 'covering knee': 1,42 'short': 1,43 'cropped': 2,44 'three-quarter': 2,45 'long': 3,46 'slack': 3,47 'of long length': 348}49socks_length_text = [50 'socks', 'stocking', 'pantyhose', 'leggings', 'sheer hosiery'51]52socks_length_attr = {53 'socks': 0,54 'stocking': 1,55 'pantyhose': 1,56 'leggings': 1,57 'sheer hosiery': 158}59hat_text = ['hat', 'cap', 'chapeau']60eyeglasses_text = ['sunglasses']61belt_text = ['belt', 'with a dress tied around the waist']62outer_shape_text = [63 'with outer clothing open', 'with outer clothing unzipped',64 'covering inner clothes', 'with outer clothing zipped'65]66outer_shape_attr = {67 'with outer clothing open': 0,68 'with outer clothing unzipped': 0,69 'covering inner clothes': 1,70 'with outer clothing zipped': 171}72 73upper_types = [74 'T-shirt', 'shirt', 'sweater', 'hoodie', 'tops', 'blouse', 'Basic Tee'75]76outer_types = [77 'jacket', 'outer clothing', 'coat', 'overcoat', 'blazer', 'outerwear',78 'duffle', 'cardigan'79]80skirt_types = ['skirt']81dress_types = ['dress']82pant_types = ['jeans', 'pants', 'trousers']83rompers_types = ['rompers', 'bodysuit', 'jumpsuit']84 85attr_names_list = [86 'gender', 'hair length', '0 upper clothing length',87 '1 lower clothing length', '2 socks', '3 hat', '4 eyeglasses', '5 belt',88 '6 opening of outer clothing', '7 upper clothes', '8 outer clothing',89 '9 skirt', '10 dress', '11 pants', '12 rompers'90]91 92 93def generate_shape_attributes(user_shape_texts):94 model = SentenceTransformer('all-MiniLM-L6-v2')95 parsed_texts = user_shape_texts.split(',')96 97 text_num = len(parsed_texts)98 99 human_attr = [0, 0]100 attr = [1, 3, 0, 0, 0, 3, 1, 1, 0, 0, 0, 0, 0]101 102 changed = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]103 for text_id, text in enumerate(parsed_texts):104 user_embeddings = model.encode(text)105 if ('man' in text) and (text_id == 0):106 human_attr[0] = 0107 human_attr[1] = 0108 109 if ('woman' in text or 'lady' in text) and (text_id == 0):110 human_attr[0] = 1111 human_attr[1] = 2112 113 if (not changed[0]) and (text_id == 1):114 # upper length115 predefined_embeddings = model.encode(upper_length_text)116 similarities = util.dot_score(user_embeddings,117 predefined_embeddings)118 arg_idx = torch.argmax(similarities).item()119 attr[0] = upper_length_attr[upper_length_text[arg_idx]]120 changed[0] = 1121 122 if (not changed[1]) and ((text_num == 2 and text_id == 1) or123 (text_num > 2 and text_id == 2)):124 # lower length125 predefined_embeddings = model.encode(lower_length_text)126 similarities = util.dot_score(user_embeddings,127 predefined_embeddings)128 arg_idx = torch.argmax(similarities).item()129 attr[1] = lower_length_attr[lower_length_text[arg_idx]]130 changed[1] = 1131 132 if (not changed[2]) and (text_id > 2):133 # socks length134 predefined_embeddings = model.encode(socks_length_text)135 similarities = util.dot_score(user_embeddings,136 predefined_embeddings)137 arg_idx = torch.argmax(similarities).item()138 if similarities[0][arg_idx] > 0.7:139 attr[2] = arg_idx + 1140 changed[2] = 1141 142 if (not changed[3]) and (text_id > 2):143 # hat144 predefined_embeddings = model.encode(hat_text)145 similarities = util.dot_score(user_embeddings,146 predefined_embeddings)147 if similarities[0][0] > 0.7:148 attr[3] = 1149 changed[3] = 1150 151 if (not changed[4]) and (text_id > 2):152 # glasses153 predefined_embeddings = model.encode(eyeglasses_text)154 similarities = util.dot_score(user_embeddings,155 predefined_embeddings)156 arg_idx = torch.argmax(similarities).item()157 if similarities[0][arg_idx] > 0.7:158 attr[4] = arg_idx + 1159 changed[4] = 1160 161 if (not changed[5]) and (text_id > 2):162 # belt163 predefined_embeddings = model.encode(belt_text)164 similarities = util.dot_score(user_embeddings,165 predefined_embeddings)166 arg_idx = torch.argmax(similarities).item()167 if similarities[0][arg_idx] > 0.7:168 attr[5] = arg_idx + 1169 changed[5] = 1170 171 if (not changed[6]) and (text_id == 3):172 # outer coverage173 predefined_embeddings = model.encode(outer_shape_text)174 similarities = util.dot_score(user_embeddings,175 predefined_embeddings)176 arg_idx = torch.argmax(similarities).item()177 if similarities[0][arg_idx] > 0.7:178 attr[6] = arg_idx179 changed[6] = 1180 181 if (not changed[10]) and (text_num == 2 and text_id == 1):182 # dress_types183 predefined_embeddings = model.encode(dress_types)184 similarities = util.dot_score(user_embeddings,185 predefined_embeddings)186 similarity_skirt = util.dot_score(user_embeddings,187 model.encode(skirt_types))188 if similarities[0][0] > 0.5 and similarities[0][189 0] > similarity_skirt[0][0]:190 attr[10] = 1191 attr[7] = 0192 attr[8] = 0193 attr[9] = 0194 attr[11] = 0195 attr[12] = 0196 197 changed[0] = 1198 changed[10] = 1199 changed[7] = 1200 changed[8] = 1201 changed[9] = 1202 changed[11] = 1203 changed[12] = 1204 205 if (not changed[12]) and (text_num == 2 and text_id == 1):206 # rompers_types207 predefined_embeddings = model.encode(rompers_types)208 similarities = util.dot_score(user_embeddings,209 predefined_embeddings)210 max_similarity = torch.max(similarities).item()211 if max_similarity > 0.6:212 attr[12] = 1213 attr[7] = 0214 attr[8] = 0215 attr[9] = 0216 attr[10] = 0217 attr[11] = 0218 219 changed[12] = 1220 changed[7] = 1221 changed[8] = 1222 changed[9] = 1223 changed[10] = 1224 changed[11] = 1225 226 if (not changed[7]) and (text_num > 2 and text_id == 1):227 # upper_types228 predefined_embeddings = model.encode(upper_types)229 similarities = util.dot_score(user_embeddings,230 predefined_embeddings)231 max_similarity = torch.max(similarities).item()232 if max_similarity > 0.6:233 attr[7] = 1234 changed[7] = 1235 236 if (not changed[8]) and (text_id == 3):237 # outer_types238 predefined_embeddings = model.encode(outer_types)239 similarities = util.dot_score(user_embeddings,240 predefined_embeddings)241 arg_idx = torch.argmax(similarities).item()242 if similarities[0][arg_idx] > 0.7:243 attr[6] = outer_shape_attr[outer_shape_text[arg_idx]]244 attr[8] = 1245 changed[8] = 1246 247 if (not changed[9]) and (text_num > 2 and text_id == 2):248 # skirt_types249 predefined_embeddings = model.encode(skirt_types)250 similarity_skirt = util.dot_score(user_embeddings,251 predefined_embeddings)252 similarity_dress = util.dot_score(user_embeddings,253 model.encode(dress_types))254 if similarity_skirt[0][0] > 0.7 and similarity_skirt[0][255 0] > similarity_dress[0][0]:256 attr[9] = 1257 attr[10] = 0258 changed[9] = 1259 changed[10] = 1260 261 if (not changed[11]) and (text_num > 2 and text_id == 2):262 # pant_types263 predefined_embeddings = model.encode(pant_types)264 similarities = util.dot_score(user_embeddings,265 predefined_embeddings)266 max_similarity = torch.max(similarities).item()267 if max_similarity > 0.6:268 attr[11] = 1269 attr[9] = 0270 attr[10] = 0271 attr[12] = 0272 changed[11] = 1273 changed[9] = 1274 changed[10] = 1275 changed[12] = 1276 277 return human_attr + attr278 279 280def generate_texture_attributes(user_text):281 parsed_texts = user_text.split(',')282 283 attr = []284 for text in parsed_texts:285 if ('pure color' in text) or ('solid color' in text):286 attr.append(4)287 elif ('spline' in text) or ('stripe' in text):288 attr.append(3)289 elif ('plaid' in text) or ('lattice' in text):290 attr.append(5)291 elif 'floral' in text:292 attr.append(1)293 elif 'denim' in text:294 attr.append(0)295 else:296 attr.append(17)297 298 if len(attr) == 1:299 attr.append(attr[0])300 attr.append(17)301 302 if len(attr) == 2:303 attr.append(17)304 305 return attr306 307 308if __name__ == "__main__":309 user_request = input('Enter your request: ')310 while user_request != '\\q':311 attr = generate_shape_attributes(user_request)312 print(attr)313 for attr_name, attr_value in zip(attr_names_list, attr):314 print(attr_name, attr_value)315 user_request = input('Enter your request: ')316