File size: 2,390 Bytes
3cc129a |
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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import gradio as gr
import cv2
import numpy as np
from PIL import Image
from ultralytics import YOLO
# Load the YOLO models
model_all = YOLO('best.pt') # Model for helmet, license plate, and motorcyclist
model_np = YOLO('best_1.pt') # Model for number plate detection
def predict(image):
# Convert PIL Image to numpy array
img_array = np.array(image)
# Perform inference with both models
results_all = model_all(img_array)
results_np = model_np(img_array)
# Plot results from the first model (all detections)
for r in results_all:
im_array = r.plot()
im = Image.fromarray(im_array[..., ::-1]) # RGB PIL image
# Initialize counters and confidence lists
class_counts_all = {i: 0 for i in range(len(model_all.names))}
class_confidences_all = {i: [] for i in range(len(model_all.names))}
np_count = 0
np_confidences = []
# Process results from the first model (all detections)
for box in results_all[0].boxes:
cls = int(box.cls[0])
conf = float(box.conf[0])
class_counts_all[cls] += 1
class_confidences_all[cls].append(conf)
# Process results from the second model (number plate detection)
for box in results_np[0].boxes:
np_count += 1
np_confidences.append(float(box.conf[0]))
# Create output string
output = "Detection Results:\n"
output += "Model 1 (Helmet, License Plate, Motorcyclist):\n"
for i in range(len(model_all.names)):
count = class_counts_all[i]
avg_conf = np.mean(class_confidences_all[i]) if class_confidences_all[i] else 0
output += f"{model_all.names[i]}: {count} detections (Avg. Confidence: {avg_conf:.2f})\n"
output += "\nModel 2 (Number Plate Detection):\n"
avg_np_conf = np.mean(np_confidences) if np_confidences else 0
output += f"Number Plates: {np_count} detections (Avg. Confidence: {avg_np_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, License Plate, and Motorcyclist Detection",
description="Upload an image to detect helmets, license plates, and motorcyclists using two specialized models."
)
# Launch the interface
iface.launch(share='True') |