Update app.py

app.py

@@ -1,22 +1,74 @@
-import gradio as gr
-from transformers import pipeline
+Hugging Face's logo
+Search models, datasets, users...
+
+Spaces:
 
-# Load the DETR object detection pipeline
-object_detection_pipeline = pipeline("object-detection", model="facebook/detr-resnet-50")
+Epitech
+/
+Object-Detection 
+
+like
+1
+App
+Files
+Community
+Object-Detection
+/
+app.py
+paulmondon
+Add requirements.txt
+a72c3ec
+raw
+history
+blame
+contribute
+delete
+1.6 kB
+from transformers import DetrImageProcessor, DetrForObjectDetection
+import torch
+from PIL import Image, ImageDraw
+import gradio as gr
+import requests
+import random
 
-# Function to perform object detection on an image
 def detect_objects(image):
-    # Perform object detection
-    results = object_detection_pipeline(image)
-    # Extract bounding boxes and object labels
-    bounding_boxes = [obj["bbox"] for obj in results]
-    labels = [obj["label"] for obj in results]
-    # Return bounding boxes and labels
-    return bounding_boxes, labels
-
-# Define Gradio interface
-inputs = gr.inputs.Image()
-outputs = gr.outputs.ObjectDetection(labels=["person", "car", "truck", "bicycle", "motorcycle"]) # Customize labels as needed
-
-# Create Gradio interface
-gr.Interface(fn=detect_objects, inputs=inputs, outputs=outputs, title="Object Detection", description="Detect objects in images.").launch()
+    # Load the pre-trained DETR model
+    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # convert outputs (bounding boxes and class logits) to COCO API
+    # let's only keep detections with score > 0.9
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+
+    # Draw bounding boxes and labels on the image
+    draw = ImageDraw.Draw(image)
+    for i, (score, label, box) in enumerate(zip(results["scores"], results["labels"], results["boxes"])):
+        box = [round(i, 2) for i in box.tolist()]
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+        draw.rectangle(box, outline=color, width=3)
+        label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 2)}"
+        draw.text((box[0], box[1]), label_text, fill=color)
+
+    return image
+
+
+def upload_image(file):
+    image = Image.open(file.name)
+    image_with_boxes = detect_objects(image)
+    return image_with_boxes
+
+iface = gr.Interface(
+    fn=upload_image,
+    inputs="file",
+    outputs="image",
+    title="Object Detection",
+    description="Upload an image and detect objects using DETR model.",
+    allow_flagging=False
+)
+
+iface.launch()
+