Unique colors

app.py

@@ -8,9 +8,16 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 dino_processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-base")
 dino_model = AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-base").to(device)
 
+def generate_colors(labels):
+    import random
+    random.seed(42)
+    colors = {}
+    for label in labels:
+        colors[label] = "#{:06x}".format(random.randint(0, 0xFFFFFF))
+    return colors
+
 @spaces.GPU
 def infer(img, text_queries, score_threshold, model):
-  
     if model == "dino":
         queries = ""
         for query in text_queries:
@@ -36,9 +43,9 @@ def infer(img, text_queries, score_threshold, model):
             if score < score_threshold:
                 continue
 
-            if model == "dino":
-                if label != "":
-                    result_labels.append((box, label))
+            # Only include the labels that are part of the input text_queries
+            if model == "dino" and label in text_queries:
+                result_labels.append((box, label))
         return result_labels
 
 def query_image(img, text_queries, dino_threshold):
@@ -47,8 +54,9 @@ def query_image(img, text_queries, dino_threshold):
     annotations = []
     for box, label in dino_output:
         annotations.append({"label": label, "coordinates": {"x": box[0], "y": box[1], "width": box[2] - box[0], "height": box[3] - box[1]}})
-    return (img, {"boxes": annotations})
-
+    
+    colors = generate_colors(text_queries)
+    return (img, {"boxes": annotations, "colors": colors})
 
 dino_threshold = gr.Slider(0, 1, value=0.12, label="Grounding DINO Threshold")
 dino_output = gr.AnnotatedImage(label="Grounding DINO Output")