Upload VITON-HD/extract.py with huggingface_hub

VITON-HD/extract.py

@@ -0,0 +1,216 @@
+# extract the key phrases of garments from the sentences
+# data_path = '/mnt/workspace/workgroup/sihui.jsh/VITON-HD/try'
+# data_path = '/mnt/workspace/workgroup/sihui.jsh/humanparsing'
+data_path = '/mnt/workspace/workgroup/sihui.jsh/SHHQ'
+# data_path = '/mnt/workspace/workgroup/sihui.jsh/dresscode/dresses'
+# data_path = '/mnt/workspace/workgroup/sihui.jsh/deepfashion'
+
+# define the key words
+from tqdm import tqdm
+import os
+colors = ['black', 'white', 'grey', 'khaki', 'burgundy']
+colors = ['red', 'blue', 'green', 'yellow', 'orange', 'purple',
+          'pink', 'black', 'white', 'brown', 'gray', 'teal', 'navy', 'maroon',
+          'olive', 'cyan', 'magenta', 'beige', 'turquoise', 'lavender', 'coral', 'peach',
+          'gold', 'silver', 'indigo', 'violet', 'chartreuse', 'khaki', 'salmon', 'auburn',
+          'plum', 'ivory', 'taupe', 'fuchsia', 'cerulean', 'tan', 'russet', 'sapphire',
+          'ruby', 'emerald', 'amber', 'burgundy', 'charcoal', 'taupe', 'mauve', 'periwinkle',
+          'mustard', 'vermillion', 'cobalt', 'crimson', 'lilac', 'jade', 'khaki', 'grey']
+
+upper = ['T-shirt', 'T-shirts', 't-shirt', 'shirt', 'tee', 'top', 'jacket', 'sweatshirt', 'coat', 'suit',
+         'blouse', 'hoodie', 'sweater', 'cardigan', 'vest', 'tunic', 'dress', 'outfit', 'bodysuit', 'TEE', 'T恤', 'gown']
+
+pants = ['jeans', 'trousers', 'pants', 'shorts', 'skirt',
+         'leggings', 'sweatpants', 'joggers', 'jogger']
+
+shoes = ['sneakers', 'boots', 'sandals', 'shoes', 'shoe', 'heels', 'flats', 'loafers', 'oxfords',
+         'moccasins', 'slippers', 'wedges', 'clogs', 'espadrilles',
+         'pumps', 'platforms', 'slides', 'brogues', 'flip-flops', 'booties',
+         'stilettos', 'derbies', 'monks', 'chukkas', 'slingbacks', 'mules', 'maryjanes', 'slingbacks', 'trainers']
+with open(os.path.join(data_path, 'data.txt'), 'r') as file1,open(os.path.join(data_path,'data_chunk.txt'), 'a') as file2:
+    texts = file1.readlines()
+    pbar = tqdm(enumerate(texts),total=len(texts))
+    for cnt, line in pbar:
+        # if cnt==10:
+        #     break
+        id, sentence = line.split('\t')
+        items = {'upper': '*', 'pants': '*', 'shoes': '*'}
+
+        # if id != 'MEN_Denim_id_00000080_01_2_side.png':
+        #     continue
+        # extract the upper garments phrases
+        for word in upper:
+            if word in sentence:
+                # get the position of the word
+                pos = sentence.find(word)
+                # starting from the word, find the first word 'a' or 'an' before it
+                start_A = sentence.rfind('A ', 0, pos)
+                start_a = sentence.rfind(' a ', 0, pos)
+                start_an = sentence.rfind(' an ', 0, pos)
+                start = max(start_A, start_a, start_an)
+                # collect the phrase
+                if start != -1:
+                    # find if these are any comma in between
+                    # comma = sentence.find(',', start, pos)
+                    # if comma != -1:
+                    #     phrase = sentence[comma+1:pos+len(word)]
+                    # else:
+                    phrase = sentence[start:pos+len(word)]
+                    # filter out extra words
+                    wearing = phrase.find('wearing')
+                    if wearing != -1:
+                        phrase = phrase[wearing+8:]
+                    man = phrase.find('man')
+                    if man != -1:
+                        phrase = phrase[man+4:]
+                    woman = phrase.find('woman')
+                    if woman != -1:
+                        phrase = phrase[woman+6:]
+                    phrase = phrase.strip()
+                    # if start with 'with' or 'in', remove them
+                    if phrase.startswith('with'):
+                        phrase = phrase[4:]
+                    elif phrase.startswith('in'):
+                        phrase = phrase[2:]
+                    phrase = phrase.strip()
+                else:
+                    phrase = word
+
+                # eliminate the extra spaces
+                phrase = phrase.strip()
+                items['upper'] = phrase
+                break
+
+        for word in pants:
+            if word in sentence:
+                pos = sentence.find(word)
+                start_A = sentence.rfind('A ', 0, pos)
+                start_a = sentence.rfind(' a ', 0, pos)
+                start_an = sentence.rfind(' an ', 0, pos)
+                start = max(start_A, start_a, start_an)
+                if start != -1:
+                    # filter out the comma
+                    comma = sentence.rfind(',', start, pos)
+                    if comma != -1:
+                        phrase = sentence[comma+1:pos+len(word)]
+                    else:
+                        phrase = sentence[start:pos+len(word)]
+                    # filter out extra words
+                    wearing = phrase.find('wearing')
+                    if wearing != -1:
+                        phrase = phrase[wearing+8:]
+                    man = phrase.find('man')
+                    if man != -1:
+                        phrase = phrase[man+4:]
+                    woman = phrase.find('woman')
+                    if woman != -1:
+                        phrase = phrase[woman+6:]
+                    phrase = phrase.strip()
+                    # if start with 'with' or 'in', remove them
+                    if phrase.startswith('with'):
+                        phrase = phrase[4:]
+                    elif phrase.startswith('in'):
+                        phrase = phrase[2:]
+                    phrase = phrase.strip()
+                    # filter out upper garments
+                    for word in upper:
+                        upper_pos = phrase.find(word)
+                        if upper_pos != -1:
+                            phrase = phrase[upper_pos+len(word):]
+                            break
+                    phrase = phrase.strip()
+                    # if start with 'and', remove it
+                    if phrase.startswith('and'):
+                        phrase = phrase[3:]
+                    phrase = phrase.strip()
+                else:
+                    phrase = word
+                    # look back to see if there is any abjective
+                    for color in colors:
+                        color_pos = sentence.rfind(color, 0, pos)
+                        if color_pos != -1:
+                            phrase = sentence[color_pos:pos+len(word)]
+                            break
+                    # print(id, phrase)
+
+                phrase = phrase.strip()
+                items['pants'] = phrase
+                break
+
+        # like pants
+        for word in shoes:
+            if word in sentence:
+                pos = sentence.find(word)
+                start_A = sentence.rfind('A ', 0, pos)
+                start_a = sentence.rfind(' a ', 0, pos)
+                start_an = sentence.rfind(' an ', 0, pos)
+                start = max(start_A, start_a, start_an)
+                if start != -1:
+                    comma = sentence.rfind(',', start, pos)
+                    if comma != -1:
+                        phrase = sentence[comma+1:pos+len(word)]
+                    else:
+                        phrase = sentence[start:pos+len(word)]
+                    phrase = phrase.strip()
+                    # filter out extra words
+                    wearing = phrase.find('wearing')
+                    if wearing != -1:
+                        phrase = phrase[wearing+8:]
+                    wears = phrase.find('wears')
+                    if wears != -1:
+                        phrase = phrase[wears+6:]
+                    man = phrase.find('man')
+                    if man != -1:
+                        phrase = phrase[man+4:]
+                    woman = phrase.find('woman')
+                    if woman != -1:
+                        phrase = phrase[woman+6:]
+                    phrase = phrase.strip()
+                    # if start with 'with' or 'in', remove them
+                    if phrase.startswith('with'):
+                        phrase = phrase[4:]
+                    elif phrase.startswith('in'):
+                        phrase = phrase[2:]
+                    phrase = phrase.strip()
+                    # filter out upper garments
+                    for word in upper:
+                        upper_pos = phrase.find(word)
+                        if upper_pos != -1:
+                            phrase = phrase[upper_pos+len(word):]
+                            break
+                    # filter out pants
+                    for word in pants:
+                        pants_pos = phrase.find(word)
+                        if pants_pos != -1:
+                            phrase = phrase[pants_pos+len(word):]
+                            break
+                    phrase = phrase.strip()
+                    # if start with 'and', remove it
+                    if phrase.startswith('and'):
+                        phrase = phrase[3:]
+                    phrase = phrase.strip()
+                else:
+                    phrase = word
+                    # look back to see if there is any abjective
+                    for color in colors:
+                        color_pos = sentence.rfind(color, 0, pos)
+                        if color_pos != -1:
+                            phrase = sentence[color_pos:pos+len(word)]
+                            break
+                    # print(id, phrase)
+
+                phrase = phrase.strip()
+                items['shoes'] = phrase
+                # print(phrase)
+                break
+            
+        # flag = False
+        # for key in items:
+        #     if items[key] != '':
+        #         flag = True
+        # if not flag:
+        #     print(id)
+        file2.write(id + '\t' + sentence.strip() +'#'+ items['upper']+'#'+ items['pants']+'#'+ items['shoes']+ '\n')
+
+        # print(id, items)
+        # break