Spaces:
Running
Running
File size: 1,891 Bytes
5934d13 d9864a1 5934d13 d9864a1 5934d13 d9864a1 5934d13 d9864a1 5934d13 d9864a1 5934d13 d9864a1 5934d13 90732ab 5934d13 d9864a1 5934d13 d9864a1 5934d13 d9864a1 5934d13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import cv2
import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np
# Load the YOLOv5 model (pre-trained on COCO dataset)
model = YOLO('yolov8n.pt') # You can replace with your custom model if available
def detect_objects(image):
"""
Detect suspicious objects in the image using YOLO.
"""
# Convert PIL image to OpenCV format (numpy array)
image_np = np.array(image)
results = model(image_np) # Perform detection
# Draw bounding boxes on the detected objects
for result in results:
boxes = result.boxes # Bounding boxes
for box in boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0]) # Extract box coordinates
label = box.cls[0] # Class label
confidence = box.conf[0] # Confidence score
# Draw rectangle and label on the image
cv2.rectangle(image_np, (x1, y1), (x2, y2), (0, 255, 0), 2)
text = f"{model.names[int(label)]} ({confidence:.2f})"
cv2.putText(image_np, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
# Convert back to PIL image
processed_image = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
return processed_image
# Create a Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Suspicious Object Detection")
gr.Markdown("Upload an image or use your webcam to capture one. The app will detect objects using YOLOv5.")
# Input section
image_input = gr.Image(type="pil", label="Upload or Capture Image")
# Output section
output_image = gr.Image(type="pil", label="Processed Image with Annotations")
# Button for detection
detect_button = gr.Button("Detect Suspicious Objects")
# Link the button to the detection function
detect_button.click(detect_objects, inputs=[image_input], outputs=[output_image])
demo.launch()
|