File size: 1,577 Bytes
f35bee1 |
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 54 |
import gradio as gr
import cv2
import numpy as np
from PIL import Image
from ultralytics import YOLO
# Load the YOLO model
model = YOLO('best.pt')
def predict(image):
# Convert PIL Image to numpy array
img_array = np.array(image)
# Perform inference
results = model(img_array)
# Plot results
for r in results:
im_array = r.plot()
im = Image.fromarray(im_array[..., ::-1]) # RGB PIL image
# Get detection counts and confidences
class_counts = {i: 0 for i in range(len(model.names))}
class_confidences = {i: [] for i in range(len(model.names))}
for box in results[0].boxes:
cls = int(box.cls[0])
conf = float(box.conf[0])
class_counts[cls] += 1
class_confidences[cls].append(conf)
# Create output string
output = "Detection Results:\n"
for i in range(len(model.names)):
count = class_counts[i]
avg_conf = np.mean(class_confidences[i]) if class_confidences[i] else 0
output += f"{model.names[i]}: {count} detections (Avg. Confidence: {avg_conf:.2f})\n"
return im, output
# Create Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=[
gr.Image(type="pil", label="Detected Image"),
gr.Textbox(label="Detection Results")
],
title="Helmet and License Plate Detection",
description="Upload an image to detect helmets, license plates, and motorcyclists."
)
# Launch the interface
iface.launch() |