Update app.py

app.py

@@ -101,26 +101,43 @@ def shot(input, category, level):
 #     subColour = responses[0][0]['label'].split(" clothing:")[0]
 
 #     return subColour, mainColour, responses[0][0]['score']
-@spaces.GPU  
-def get_colour(image_urls, category):
-    colourLabels = list(COLOURS_DICT.keys())
-    for i in range(len(colourLabels)):
-        colourLabels[i] = colourLabels[i] + " clothing: " + category
 
+@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
 
-    responses = pipe(image_urls, candidate_labels=colourLabels)
-    mainColour = responses[0][0]['label'].split(" clothing:")[0]
+    # 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
 
-    labels = COLOURS_DICT[mainColour]
-    for i in range(len(labels)):
-        labels[i] = labels[i] + " clothing: " + category
-
+    # 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]
 
@@ -128,6 +145,7 @@ def get_colour(image_urls, category):
 
 
 
+
 @spaces.GPU  
 def get_predicted_attributes(image_urls, category):
     # Assuming ATTRIBUTES_DICT and pipe are defined outside this function