Update app.py

app.py

@@ -104,30 +104,27 @@ def get_colour(image_urls, category):
 
 @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())
-    
+    # Assuming ATTRIBUTES_DICT and pipe are defined outside this function
+    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 = ATTRIBUTES_DICT.get(category,{}).get(attribute,[])
+        values = ATTRIBUTES_DICT.get(category, {}).get(attribute, [])
 
         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]
+        attribute_formatted = attribute.replace("colartype", "collar").replace("sleevelength", "sleeve length").replace("fabricstyle", "fabric")
+        values_formatted = [f"{attribute_formatted}: {value}, clothing: {category}" for value in values]
 
         # Get the predicted values for the attribute
-        responses = pipe(image_urls, candidate_labels=values)
+        responses = pipe(image_urls, candidate_labels=values_formatted)
         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":
+        if attribute_formatted == "details":
             result += [response[1]['label'].split(", clothing:")[0] for response in responses]
             common_result.append(Counter(result).most_common(2))
         else:
@@ -143,12 +140,19 @@ def get_predicted_attributes(image_urls, category):
     for item in common_result:
         # Split by ': ' to separate the key and value
         key, value = item.split(': ', 1)
-        # Add to the dictionary
-        result[key] = value
+        
+        if key == "details":
+            details_split = value.split(" , ")
+            if len(details_split) == 2:
+                result["details_1"] = details_split[0]
+                result["details_2"] = details_split[1]
+            else:
+                result["details_1"] = value  # If there's only one detail, assign it to details 1
+        else:
+            result[key] = value
 
     return result
 
-
 def get_openAI_tags(image_urls):
     # Create list containing JSONs of each image URL
     imageList = []