Update app.py

app.py

@@ -108,56 +108,93 @@ def shot(input, category, level):
 
 #     return subColour, mainColour, responses[0][0]['score']
 
+# @spaces.GPU
+# def get_colour(image_urls, category):
+#     # Prepare color labels
+#     colourLabels = [f"{color} clothing: {category}" for color in COLOURS_DICT.keys()]
+#     print("Colour Labels:", colourLabels)  # Debug: Print colour labels
+#     print("Image URLs:", image_urls)       # Debug: Print image URLs
+
+#     # Split labels into two batches
+#     mid_index = len(colourLabels) // 2
+#     first_batch = colourLabels[:mid_index]
+#     second_batch = colourLabels[mid_index:]
+
+#     # Process the first batch
+#     responses_first_batch = pipe(image_urls, candidate_labels=first_batch)
+#     # Get the top 3 from the first batch
+#     top3_first_batch = sorted(responses_first_batch[0], key=lambda x: x['score'], reverse=True)[:3]
+
+#     # Process the second batch
+#     responses_second_batch = pipe(image_urls, candidate_labels=second_batch)
+#     # Get the top 3 from the second batch
+#     top3_second_batch = sorted(responses_second_batch[0], key=lambda x: x['score'], reverse=True)[:3]
+
+#     # Combine the top 3 from each batch
+#     combined_top6 = top3_first_batch + top3_second_batch
+#     # Get the final top 3 from the combined list
+#     final_top3 = sorted(combined_top6, key=lambda x: x['score'], reverse=True)[:3]
+
+#     mainColour = final_top3[0]['label'].split(" clothing:")[0]
+
+#     if mainColour not in COLOURS_DICT:
+#         return None, None, None
+
+#     # Get sub-colors for the main color
+#     labels = [f"{label} clothing: {category}" for label in COLOURS_DICT[mainColour]]
+#     print("Labels for pipe:", labels)  # Debug: Confirm labels are correct
+
+#     responses = pipe(image_urls, candidate_labels=labels)
+#     subColour = responses[0][0]['label'].split(" clothing:")[0]
+
+#     return subColour, mainColour, responses[0][0]['score']
+
+
 @spaces.GPU
 def get_colour(image_urls, category):
-    # Prepare color labels
-    colourLabels = [f"{color} clothing: {category}" for color in COLOURS_DICT.keys()]
-    print("Colour Labels:", colourLabels)  # Debug: Print colour labels
-    print("Image URLs:", image_urls)       # Debug: Print image URLs
-
-    # Split labels into two batches
-    mid_index = len(colourLabels) // 2
-    first_batch = colourLabels[:mid_index]
-    second_batch = colourLabels[mid_index:]
-
-    # Process the first batch
-    responses_first_batch = pipe(image_urls, candidate_labels=first_batch)
-    # Get the top 3 from the first batch
-    top3_first_batch = sorted(responses_first_batch[0], key=lambda x: x['score'], reverse=True)[:3]
-
-    # Process the second batch
-    responses_second_batch = pipe(image_urls, candidate_labels=second_batch)
-    # Get the top 3 from the second batch
-    top3_second_batch = sorted(responses_second_batch[0], key=lambda x: x['score'], reverse=True)[:3]
-
-    # Combine the top 3 from each batch
-    combined_top6 = top3_first_batch + top3_second_batch
-    # Get the final top 3 from the combined list
-    final_top3 = sorted(combined_top6, key=lambda x: x['score'], reverse=True)[:3]
-
-    mainColour = final_top3[0]['label'].split(" clothing:")[0]
-
-    if mainColour not in COLOURS_DICT:
+    colour_labels = [f"{colour}, clothing: {category}" for colour in COLOURS_DICT.keys()]
+    responses = pipe(image_urls, candidate_labels=colour_labels, device=device)
+    
+    main_colour, main_score = get_most_common_label(responses)
+    if main_colour not in COLOURS_DICT:
         return None, None, None
 
-    # Get sub-colors for the main color
-    labels = [f"{label} clothing: {category}" for label in COLOURS_DICT[mainColour]]
-    print("Labels for pipe:", labels)  # Debug: Confirm labels are correct
+    score = [main_score]
 
-    responses = pipe(image_urls, candidate_labels=labels)
-    subColour = responses[0][0]['label'].split(" clothing:")[0]
+    labels = COLOURS_DICT[main_colour]
+    if main_colour == "multicolor":
+        labels = [label for key, values in COLOURS_DICT.items() if key != main_colour for label in values]
 
-    return subColour, mainColour, responses[0][0]['score']
+    labels = [f"{label}, clothing: {category}" for label in labels]
+    responses = pipe(image_urls, candidate_labels=labels, device=device)
+    
+    most_common, sub_score = get_most_common_label(responses)
+    sub_colours = [most_common]
+    score.append(sub_score)
+
+    if main_colour == "multicolor":
+        sub_key = next(key for key, values in COLOURS_DICT.items() if most_common in values)
+        labels = [label for key, values in COLOURS_DICT.items() if key not in {main_colour, sub_key} for label in values]
+        labels = [f"{label}, clothing: {category}" for label in labels]
+        responses = pipe(image_urls, candidate_labels=labels, device=device)
+        
+        most_common, tertiary_score = get_most_common_label(responses)
+        sub_colours.append(most_common)
+        score.append(tertiary_score)
+    
+    return sub_colours, main_colour, score
 
 
 
-# Function for get_predicted_attributes
+# Function for get_predicted_attributes and get_colour
 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)]
+        
+    max_label = max(feature_scores, key=feature_scores.get)
+return max_label, feature_scores[max_label] / len(responses)
 
 @spaces.GPU  
 def get_predicted_attributes(image_urls, category):