Update app.py

app.py

@@ -3,11 +3,14 @@ import torch
 from PIL import Image
 from torchvision.transforms import functional as F
 from transformers import DetrImageProcessor, DetrForObjectDetection
+import cv2
+import numpy as np
 
 # Load the pretrained DETR model
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
 
+# Define the object detection function
 def detect_objects(frame):
     # Convert the frame to PIL image
     image = Image.fromarray(frame)
@@ -25,16 +28,17 @@ def detect_objects(frame):
     # Draw bounding boxes on the frame
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
         box = [round(i, 2) for i in box.tolist()]
-        frame = gr.draw_box(frame, box, label=model.config.id2label[label.item()], color=(0, 255, 0))
+        frame = cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)
+        frame = cv2.putText(frame, f'{model.config.id2label[label.item()]}: {round(score.item(), 3)}',
+                            (int(box[0]), int(box[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
 
-    # Convert frame back to numpy array for Gradio
-    return np.array(frame)
+    return frame
 
 # Define the Gradio interface
 iface = gr.Interface(
     fn=detect_objects,
     inputs=gr.Video(),
-    outputs="video",
+    outputs="numpy_image",
     live=True,
 )