Update app.py

app.py

@@ -10,48 +10,32 @@ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revisi
 model.eval()
 
 # Function for live object detection from the camera
-def live_object_detection():
-    # Open a connection to the camera (replace with your own camera setup)
-    cap = cv2.VideoCapture(0)
+def live_object_detection(image_pil):
+    # Convert the frame to PIL Image
+    frame_pil = Image.fromarray(cv2.cvtColor(image_pil, cv2.COLOR_BGR2RGB))
 
-    while True:
-        # Capture frame-by-frame
-        ret, frame = cap.read()
+    # Process the frame with the DETR model
+    inputs = processor(images=frame_pil, return_tensors="pt")
+    outputs = model(**inputs)
 
-        # Convert the frame to PIL Image
-        frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    # convert outputs (bounding boxes and class logits) to COCO API
+    # let's only keep detections with score > 0.9
+    target_sizes = torch.tensor([frame_pil.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
 
-        # Process the frame with the DETR model
-        inputs = processor(images=frame_pil, return_tensors="pt")
-        outputs = model(**inputs)
+    # Draw bounding boxes on the frame
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [int(round(i)) for i in box.tolist()]
+        cv2.rectangle(image_pil, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
+        cv2.putText(image_pil, f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}",
+                    (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
 
-        # convert outputs (bounding boxes and class logits) to COCO API
-        # let's only keep detections with score > 0.9
-        target_sizes = torch.tensor([frame_pil.size[::-1]])
-        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
-
-        # Draw bounding boxes on the frame
-        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-            box = [int(round(i)) for i in box.tolist()]
-            cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
-            cv2.putText(frame, f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}", 
-                        (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-
-        # Display the resulting frame
-        cv2.imshow('Object Detection', frame)
-
-        # Break the loop when 'q' key is pressed
-        if cv2.waitKey(1) & 0xFF == ord('q'):
-            break
-
-    # Release the camera and close all windows
-    cap.release()
-    cv2.destroyAllWindows()
+    return image_pil
 
 # Define the Gradio interface
 iface = gr.Interface(
     fn=live_object_detection,
-    inputs="webcam",
+    inputs="image",
     outputs="image",
     live=True,
 )