Update app.py

app.py

@@ -1,35 +1,67 @@
 # app.py
 
+# app.py
+
 from transformers import pipeline
 import gradio as gr
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
+import random
 
-# Load the object detection pipeline with YOLO
+# Load the YOLO-based object detection pipeline
 detector = pipeline("object-detection", model="hustvl/yolos-tiny")
 
-# Define the detection function
+# Generate a random color for each label
+label_colors = {}
+
+def get_color(label):
+    if label not in label_colors:
+        label_colors[label] = (
+            random.randint(0, 255),
+            random.randint(0, 255),
+            random.randint(0, 255)
+        )
+    return label_colors[label]
+
+# Detection function
 def detect_objects(img):
     results = detector(img)
 
-    # Draw bounding boxes
     draw = ImageDraw.Draw(img)
+    font = ImageFont.load_default()
+
     for obj in results:
+        label = obj["label"]
+        score = obj["score"]
         box = obj["box"]
-        label = f"{obj['label']} ({obj['score']:.2f})"
-        draw.rectangle([box["xmin"], box["ymin"], box["xmax"], box["ymax"]], outline="red", width=3)
-        draw.text((box["xmin"], box["ymin"] - 10), label, fill="red")
-    
+        color = get_color(label)
+
+        # Draw bounding box
+        draw.rectangle(
+            [box["xmin"], box["ymin"], box["xmax"], box["ymax"]],
+            outline=color,
+            width=3
+        )
+
+        # Prepare label with confidence
+        label_text = f"{label} ({score:.2f})"
+        text_size = draw.textsize(label_text, font=font)
+        text_background = [box["xmin"], box["ymin"] - text_size[1], box["xmin"] + text_size[0], box["ymin"]]
+
+        # Draw background for text
+        draw.rectangle(text_background, fill=color)
+        draw.text((box["xmin"], box["ymin"] - text_size[1]), label_text, fill="black", font=font)
+
     return img
 
-# Create the Gradio interface
+# Gradio interface
 interface = gr.Interface(
     fn=detect_objects,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="YOLO Object Detection",
-    description="Upload an image to detect objects using a YOLOS-Tiny model."
+    title="YOLO Object Detection with Color-coded Labels",
+    description="Upload an image. Detected objects are shown with bounding boxes and color-coded labels using YOLOS-Tiny."
 )
 
-# Launch the app
 if __name__ == "__main__":
     interface.launch()
+