Update app.py

app.py

@@ -1,23 +1,45 @@
 from transformers import DetrFeatureExtractor, DetrForObjectDetection
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 import requests
 import gradio as gr
 
+# Draw bounding box definition
+def draw_bounding_box(im, score, label, xmin, ymin, xmax, ymax, index, num_boxes):
+  """ Draw a bounding box. """
+  # Draw the actual bounding box
+  outline = 'blue'
+  im_with_rectangle = ImageDraw.Draw(im)  
+  im_with_rectangle.rounded_rectangle((xmin, ymin, xmax, ymax), outline = outline, width = 3, radius = 10)
+
+  # Return the result
+  return im
+
 def object_classify(img):
   feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50')
   model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
   
-  inputs = feature_extractor(images=img, return_tensors="pt")
-  outputs = model(**inputs)
-  # model predicts bounding boxes and corresponding COCO classes
-  logits = outputs.logits
-  bboxes = outputs.pred_boxes
+  object_detector = pipeline("object-detection", model = model, feature_extractor = feature_extractor)
+  bboxes = object_detector(img)
+
   
-  probas = logits.softmax(-1)[0, :, :-1]
-  keep = probas.max(-1).values > 0.7
-  return keep
+  # Iteration elements
+  num_boxes = len(bounding_boxes)
+  index = 0
+
+  # Draw bounding box for each result
+  for bounding_box in bounding_boxes:
+    box = bounding_box['box']
+    #Draw the bounding box
+    output_image = draw_bounding_box(img, bounding_box['score'],
+                                  bounding_box['label'],
+                                  box['xmin'], box['ymin'],
+                                  box['xmax'], box['ymax'],
+                                  index, num_boxes)
+  index += 1
+  return output_image
+
   
 interface=gr.Interface(fn=object_classify,
-                       inputs=gr.inputs.Image(shape=(224,224),label='Insert Image'),
+                       inputs=gr.inputs.Image(shape=(224,224),label='Insert Image'),outputs = gr.outputs.Image(),
                        examples = [],description='??',allow_flagging="never")
 interface.launch()