Vishaltiwari2019 commited on
Commit
b6e3dbd
1 Parent(s): 6ac27af

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +16 -33
app.py CHANGED
@@ -1,41 +1,44 @@
1
- import gradio as gr
2
  from transformers import DetrImageProcessor, DetrForObjectDetection
3
  import torch
4
- from PIL import Image, ImageDraw
 
 
5
  import random
6
 
7
  def detect_objects(image):
 
8
  processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
9
  model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
10
 
11
  inputs = processor(images=image, return_tensors="pt")
12
  outputs = model(**inputs)
13
 
 
 
14
  target_sizes = torch.tensor([image.size[::-1]])
15
  results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
16
 
 
17
  draw = ImageDraw.Draw(image)
18
- labels = []
19
  for i, (score, label, box) in enumerate(zip(results["scores"], results["labels"], results["boxes"])):
20
  box = [round(i, 2) for i in box.tolist()]
21
  color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
22
  draw.rectangle(box, outline=color, width=3)
23
  label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 2)}"
 
24
  draw.text((box[0], box[1]), label_text, fill=color,)
25
- labels.append(label_text)
 
 
26
 
27
- return image, labels
28
 
29
  def upload_image(file):
30
  image = Image.open(file.name)
31
- image_with_boxes, labels = detect_objects(image)
32
- return image_with_boxes, labels
33
 
34
- def show_labels(labels):
35
- return "\n".join(labels)
36
-
37
- # Interface to display the image with bounding boxes
38
- iface_objects = gr.Interface(
39
  fn=upload_image,
40
  inputs="file",
41
  outputs=["image", "text"],
@@ -44,24 +47,4 @@ iface_objects = gr.Interface(
44
  allow_flagging=False
45
  )
46
 
47
- # Interface to display the detected labels
48
- iface_labels = gr.Interface(
49
- fn=show_labels,
50
- inputs="text",
51
- outputs="text",
52
- title="Detected Labels",
53
- description="Displays the labels detected in the uploaded image.",
54
- allow_flagging=False
55
- )
56
-
57
- # Combine interfaces with a tapped interface
58
- interface = gr.Interface(
59
- [iface_objects, iface_labels],
60
- inputs="text",
61
- outputs="text",
62
- title="Object Detection with Labels",
63
- description="Upload an image and view detected objects and labels.",
64
- allow_flagging=False
65
- )
66
-
67
- interface.launch()
 
 
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"],
 
47
  allow_flagging=False
48
  )
49
 
50
+ iface.launch()