Update app.py

app.py

@@ -7,6 +7,18 @@ import torch
 image_processor = AutoImageProcessor.from_pretrained('hustvl/yolos-small')
 model = AutoModelForObjectDetection.from_pretrained('hustvl/yolos-small')
 
+colors = ["red",
+          "orange",
+          "yellow",
+          "green",
+          "blue",
+          "indigo",
+          "violet",
+          "brown",
+          "black",
+          "slategray",
+         ]
+
 def detect(image): 
     inputs = image_processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
@@ -22,23 +34,32 @@ def detect(image):
     # label and the count
     counts = {}
 
-    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        box = [round(i, 4) for i in box.tolist()]
+    for score, label in zip(results["scores"], results["labels"]):
         label_name = model.config.id2label[label.item()]
         if label_name not in counts:
             counts[label_name] = 0
         counts[label_name] += 1
 
-        x1, y1, x2, y2 = tuple(box)
-        draw.rectangle((x1, y1, x2, y2), outline=(128, 128, 50), width=2)
-        draw.text((x1, y1), label_name, fill="white")
+    count_results = {k: v for k, v in sorted(counts.items(), key=lambda item: item[1])}[:10]
+    label2color = {}
+    for idx, label in enumerate(count_results):
+        label2color[label] = colors[idx]
+
+    for label, box in zip(results["labels"], results["boxes"]):
+        label_name = model.config.id2label[label.item()]
+
+        if label_name in count_results:
+            box = [round(i, 4) for i in box.tolist()]
+            x1, y1, x2, y2 = tuple(box)
+            draw.rectangle((x1, y1, x2, y2), outline=label2color[label_name], width=2)
+            draw.text((x1, y1), label_name, fill="white")
 
     df = pd.DataFrame({
-        'label': [label for label in counts],
-        'counts': [counts[label] for label in counts]
+        'label': [label for label in count_results],
+        'counts': [counts[label] for label in count_results]
     })
-        
-    return image, df, counts
+    
+    return image, df, count_results
 
 demo = gr.Interface(
     fn=detect,