Update app.py

app.py

@@ -3,6 +3,7 @@ import cv2
 import torch
 from PIL import Image
 from transformers import DetrImageProcessor, DetrForObjectDetection
+import numpy as np
 
 # Load the pre-trained DETR model
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
@@ -33,33 +34,54 @@ def image_object_detection(image_pil):
     return image_np
 
 # Function for live object detection from the camera
-def live_object_detection(image_pil):
-    # Process the frame with the DETR model
-    inputs = processor(images=image_pil, return_tensors="pt")
-    outputs = model(**inputs)
+def live_object_detection():
+    # Open a connection to the camera (replace with your own camera setup)
+    cap = cv2.VideoCapture(0)
 
-    # convert outputs (bounding boxes and class logits) to COCO API
-    # let's only keep detections with score > 0.9
-    target_sizes = torch.tensor([image_pil.size[::-1]])
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+    while True:
+        # Capture frame-by-frame
+        ret, frame = cap.read()
 
-    # Draw bounding boxes on the image
-    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)
-        label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}"
-        cv2.putText(image_pil, label_text, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        # Convert the frame to PIL Image
+        frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+
+        # Process the frame with the DETR model
+        inputs = processor(images=frame_pil, return_tensors="pt")
+        outputs = model(**inputs)
+
+        # 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]
 
-    return image_pil
+        # 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)
+            label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}"
+            cv2.putText(frame, label_text, (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()
 
 # Define the Gradio interface
 iface = gr.Interface(
     fn=[image_object_detection, live_object_detection],
     inputs=[
-        gr.Image(type="pil", label="Upload an image for object detection"),
-        "webcam",
+        gr.Image(type="pil", label="Upload an image for object detection") # Remove this line
+    ],
+    outputs=[
+        "image",
+        "image",
     ],
-    outputs=["image", "image"],
     live=True,
 )