Update app.py

app.py

@@ -11,7 +11,7 @@ from io import BytesIO
 import face_recognition
 from turtle import title
 from openai import OpenAI
-from collections import Counter
+from collections import Counter, defaultdict
 from transformers import pipeline
 
 import urllib.request
@@ -145,6 +145,13 @@ def get_colour(image_urls, category):
 
 
 
+# Function for get_predicted_attributes
+def get_most_common_label(responses):
+    feature_scores = defaultdict(float)
+    for response in responses:
+        label, score = response[0]['label'].split(", clothing:")[0], response[0]['score']
+        feature_scores[label] += score
+    return max(feature_scores, key=feature_scores.get), feature_scores[max(feature_scores, key=feature_scores.get)]
 
 @spaces.GPU  
 def get_predicted_attributes(image_urls, category):
@@ -159,56 +166,38 @@ def get_predicted_attributes(image_urls, category):
         if len(values) == 0:
             continue
 
-        # Adjust labels for the pipeline to be in format: "{attr}: {value}, clothing: {category}"
-        attribute_formatted = attribute.replace("colartype", "collar").replace("sleevelength", "sleeve length").replace("fabricstyle", "fabric")
-        values_formatted = [f"{attribute_formatted}: {value.strip()}, clothing: {category}" for value in values]
-
-        # Get the predicted values for the attribute
-        responses = pipe(image_urls, candidate_labels=values_formatted)
-        result = [response[0]['label'].split(", clothing:")[0] for response in responses]
-        common_result.append(Counter(result).most_common(1))
-        print("Common Result:", common_result)
-
-        # If attribute is details, then remove the obtained label from the values, and get the next most common
-        if attribute == "details":
-            print(values_formatted, result[0] + f", clothing: {category}")
-            values_formatted.remove(result[0] + f", clothing: {category}")
-            responses = pipe(image_urls, candidate_labels=values_formatted)
-            i = 0
-            print("Responses: ", responses)
-            while len(responses) > 0 and responses[0][0]['score'] > 0.6 and i < 2:
-                result = [response[0]['label'].split(", clothing:")[0] for response in responses]
-                common_result.append(Counter(result).most_common(1))
-
-                values_formatted.remove(result[0] + f", clothing: {category}")
-                responses = pipe(image_urls, candidate_labels=values_formatted)
-                print("Responses: ", responses)
-                i += 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])
-    
-    result = {}
-
-    # Iterate through the list and split each item into key and value
-    for item in common_result:
-        # Split by ': ' to separate the key and value
-        key, value = item.split(': ', 1)
-        
+       # Adjust labels for the pipeline
+    attribute = attribute.replace("colartype", "collar").replace("sleevelength", "sleeve length").replace("fabricstyle", "fabric")
+    values = [f"{attribute}: {value.strip()}, clothing: {category}" for value in values]
+
+    # Get the predicted values for the attribute
+    responses = pipe(product['Images'].values[0], candidate_labels=values, device=device)
+    most_common, score = get_most_common_label(responses)
+    common_result.append(most_common)
+
+    if attribute == "details":
+        # Process additional details labels if the score is higher than 0.8
+        for _ in range(2):
+            values = [value for value in values if value != f"{most_common}, clothing: {category}"]
+            responses = pipe(product['Images'].values[0], candidate_labels=values, device=device)
+            most_common, score = get_most_common_label(responses)
+            if score > 0.8:
+                common_result.append(most_common)
+
+    # Convert common_result into a dictionary
+    final = {}
+    details_count = 0
+
+    for result in common_result:
+        result = result.replace("collar", "colartype").replace("sleeve length", "sleevelength").replace("fabric", "fabricstyle")
+        key, value = result.split(": ")
         if key == "details":
-            print(value)
-            details_split = value.split(" , ")
-            if len(details_split) == 2:
-                result["details1"] = details_split[0].lower()
-                result["details2"] = details_split[1].lower()
-            else:
-                result["details1"] = value.lower()  # If there's only one detail, assign it to details 1
-        else:
-            result[key.lower().replace("collar", "colartype").replace("sleeve length", "sleevelength").replace("fabric", "fabricstyle")] = value.lower()
+            if details_count > 0:
+                key += str(details_count)
+            details_count += 1
+        final[key] = value.lower()
 
-    return result
+    return final
 
 def get_openAI_tags(image_urls):
     # Create list containing JSONs of each image URL