PopTeil commited on
Commit
dd76981
·
verified ·
1 Parent(s): d00b7e1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -43
app.py CHANGED
@@ -1,51 +1,26 @@
1
  from transformers import DetrImageProcessor, DetrForObjectDetection
2
  import torch
3
- from PIL import Image, ImageDraw, ImageFont # Import ImageFont
4
- import gradio as gr
5
  import requests
6
- import random
7
 
8
- def detect_objects(image):
9
- # Load the pre-trained DETR model
10
- processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
11
- model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
12
 
13
- inputs = processor(images=image, return_tensors="pt")
14
- outputs = model(**inputs)
 
15
 
16
- # convert outputs (bounding boxes and class logits) to COCO API
17
- # let's only keep detections with score > 0.9
18
- target_sizes = torch.tensor([image.size[::-1]])
19
- results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
20
 
21
- # Draw bounding boxes and labels on the image
22
- draw = ImageDraw.Draw(image)
23
- detected_objects = []
24
- for i, (score, label, box) in enumerate(zip(results["scores"], results["labels"], results["boxes"])):
25
- box = [round(i, 2) for i in box.tolist()]
26
- color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
27
- draw.rectangle(box, outline=color, width=3)
28
- label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 2)}"
29
- # Larger and bolder font
30
- draw.text((box[0], box[1]), label_text, fill=color,)
31
- detected_objects.append(model.config.id2label[label.item()])
32
 
33
- return image, ', '.join(detected_objects)
34
-
35
-
36
- def upload_image(file):
37
- image = Image.open(file.name)
38
- image_with_boxes, detected_objects = detect_objects(image)
39
- return image_with_boxes, detected_objects
40
-
41
- iface = gr.Interface(
42
- fn=upload_image,
43
- inputs="file",
44
- outputs=["image", "text"],
45
- title="Object Detection",
46
- description="Upload an image and detect objects using DETR model.",
47
- allow_flagging=False,
48
- css="style.css" # Path to your custom CSS file
49
- )
50
-
51
- iface.launch()
 
1
  from transformers import DetrImageProcessor, DetrForObjectDetection
2
  import torch
3
+ from PIL import Image
 
4
  import requests
 
5
 
6
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
7
+ image = Image.open(requests.get(url, stream=True).raw)
 
 
8
 
9
+ # you can specify the revision tag if you don't want the timm dependency
10
+ processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-101", revision="no_timm")
11
+ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-101", revision="no_timm")
12
 
13
+ inputs = processor(images=image, return_tensors="pt")
14
+ outputs = model(**inputs)
 
 
15
 
16
+ # convert outputs (bounding boxes and class logits) to COCO API
17
+ # let's only keep detections with score > 0.9
18
+ target_sizes = torch.tensor([image.size[::-1]])
19
+ results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
 
 
 
 
 
 
 
20
 
21
+ for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
22
+ box = [round(i, 2) for i in box.tolist()]
23
+ print(
24
+ f"Detected {model.config.id2label[label.item()]} with confidence "
25
+ f"{round(score.item(), 3)} at location {box}"
26
+ )