Update app.py

app.py

@@ -10,24 +10,31 @@ from transformers import pipeline
 import ast
 pipe = pipeline("zero-shot-image-classification", model="patrickjohncyh/fashion-clip")
 
-file_path = 'color_config.json'
+color_file_path = 'color_config.json'
+attributes_file_path = 'attributes_config.json'
 
-# Open and read the JSON file
-with open(file_path, 'r') as file:
-    data = json.load(file)
+# Open and read the COLOR JSON file
+with open(color_file_path, 'r') as file:
+    color_data = json.load(file)
 
-COLOURS_DICT = data['color_mapping']
+# Open and read the ATTRIBUTES JSON file
+with open(attributes_file_path, 'r') as file:
+    attributes_data = json.load(file)
 
+COLOURS_DICT = color_data['color_mapping']
+ATTRIBUTES_DICT = attributes_data['attribute_mapping']
 
 
 def shot(input, category):
     subColour,mainColour,score = get_colour(ast.literal_eval(str(input)),category)
+    common_result = get_predicted_attributes(ast.literal_eval(str(input)),category)
     return {
         "colors":{
             "main":mainColour,
             "sub":subColour,
             "score":round(score*100,2)
         }   
+        "attributes":common_result
     }
 
 
@@ -57,6 +64,43 @@ def get_colour(image_urls, category):
 
     return subColour, mainColour, responses[0][0]['score']
 
+@spaces.GPU  
+def get_predicted_attributes(image_urls, category):
+    # Get the predicted attributes for the image
+    # attributes = get_category_attributes(category)
+    
+    attributes = list(ATTRIBUTES_DICT.get(category,{}).keys())
+    
+    # Mapping of possible values per attribute
+    common_result = []
+    for attribute in attributes:
+        # values = get_attribute_values(attribute, category)
+        values = list(ATTRIBUTES_DICT.get(category,{}).get(attribute,{}).keys())
+
+        if len(values) == 0:
+            continue
+
+        # Adjust labels for the pipeline to be in format: "{attr}: {value}, clothing: {category}"
+        attribute = attribute.replace("colartype", "collar").replace("sleevelength", "sleeve length").replace("fabricstyle", "fabric")
+        values = [f"{attribute}: {value}, clothing: {category}" for value in values]
+
+        # Get the predicted values for the attribute
+        responses = pipe(image_urls, candidate_labels=values, device=device)
+        result = [response[0]['label'].split(", clothing:")[0] for response in responses]
+
+        # If attribute is details, then get the top 2 most common labels
+        if attribute == "details":
+            result += [response[1]['label'].split(", clothing:")[0] for response in responses]
+            common_result.append(Counter(result).most_common(2))
+        else:
+            common_result.append(Counter(result).most_common(1))
+
+    # Clean up the results into one long string
+    for i, result in enumerate(common_result):
+        common_result[i] = ", ".join([f"{x[0]}" for x in result])
+
+    return common_result
+