Change label

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 from transformers import DetrImageProcessor, DetrForObjectDetection
 import torch
 
@@ -8,6 +8,9 @@ model_name = "facebook/detr-resnet-50"
 processor = DetrImageProcessor.from_pretrained(model_name)
 model = DetrForObjectDetection.from_pretrained(model_name)
 
+# Load default font
+font = ImageFont.load_default()
+
 # Main function: takes an image and returns it with boxes and labels
 def detect_objects(image):
     inputs = processor(images=image, return_tensors="pt")
@@ -22,12 +25,26 @@ def detect_objects(image):
     # Draw bounding boxes and labels on a copy of the image
     image_with_boxes = image.copy()
     draw = ImageDraw.Draw(image_with_boxes)
-    
+
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
         box = [round(x, 2) for x in box.tolist()]
         draw.rectangle(box, outline="red", width=3)
+
+        # Prepare label text
         label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 2)}"
-        draw.text((box[0], box[1]), label_text, fill="white")
+
+        # Measure text size
+        text_bbox = draw.textbbox((0, 0), label_text, font=font)
+        text_width = text_bbox[2] - text_bbox[0]
+        text_height = text_bbox[3] - text_bbox[1]
+
+        # Set background rectangle for text
+        text_background = [
+            box[0], box[1] - text_height,
+            box[0] + text_width, box[1]
+        ]
+        draw.rectangle(text_background, fill="black")  # Background
+        draw.text((box[0], box[1] - text_height), label_text, fill="white", font=font)
 
     return image_with_boxes