ZENLLC commited on
Commit
5cf1972
·
verified ·
1 Parent(s): 1a5b9ed

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -0
app.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import DetrImageProcessor, DetrForObjectDetection
4
+ from PIL import Image, ImageDraw
5
+
6
+ # Load the pre-trained DETR model and processor
7
+ processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
8
+ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
9
+
10
+ def detect_objects(image: Image.Image) -> Image.Image:
11
+ try:
12
+ # Preprocess the image
13
+ inputs = processor(images=image, return_tensors="pt")
14
+ outputs = model(**inputs)
15
+
16
+ # Convert outputs to bounding boxes and labels
17
+ target_sizes = torch.tensor([image.size[::-1]])
18
+ results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
19
+
20
+ # Draw bounding boxes on the image
21
+ draw = ImageDraw.Draw(image)
22
+ for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
23
+ box = [round(i, 2) for i in box.tolist()]
24
+ label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}"
25
+ draw.rectangle(box, outline="red", width=3)
26
+ draw.text((box[0], box[1]), label_text, fill="red")
27
+ return image
28
+ except Exception as e:
29
+ print("Error during detection:", e)
30
+ return image # In a robust production system, consider returning a message or a default image
31
+
32
+ # Create a Gradio interface
33
+ iface = gr.Interface(
34
+ fn=detect_objects,
35
+ inputs=gr.Image(type="pil", label="Upload an Image"),
36
+ outputs=gr.Image(label="Detection Result"),
37
+ title="Robust Object Detection with DETR",
38
+ description="Upload an image to detect objects using a pre-trained DETR model from Hugging Face Hub."
39
+ )
40
+
41
+ if __name__ == "__main__":
42
+ iface.launch()