Spaces:
Sleeping
Sleeping
File size: 3,018 Bytes
a8effe4 8ccc1e0 834767f a8effe4 8ccc1e0 a8effe4 65542b6 8ccc1e0 a8effe4 65542b6 a8effe4 65542b6 8ccc1e0 834767f a8effe4 834767f a8effe4 834767f a8effe4 834767f a8effe4 834767f a8effe4 8ccc1e0 a8effe4 8ccc1e0 a8effe4 8ccc1e0 834767f 8ccc1e0 834767f 8ccc1e0 834767f |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
from typing import get_args
from PIL import Image
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import tempfile
import os
import subprocess
from fast_alpr import ALPR
from fast_alpr.default_detector import PlateDetectorModel
from fast_alpr.default_ocr import OcrModel
# Loading YOLO model
model = YOLO('best.pt')
img_dim = (640, 640)
# Default models for plate recognition
DETECTOR_MODELS = list(get_args(PlateDetectorModel))
OCR_MODELS = list(get_args(OcrModel))
# Put global OCR first
OCR_MODELS.remove("global-plates-mobile-vit-v2-model")
OCR_MODELS.insert(0, "global-plates-mobile-vit-v2-model")
def predict(image, conf_threshold, iou_threshold):
# Resizing
image = image.resize(img_dim)
# Convert from PIL to OpenCV format
image = np.array(image)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# Run inference
results = model.predict(
image,
conf=conf_threshold,
iou=iou_threshold,
imgsz=640
)
# Get annotated image
annotated_image = results[0].plot()
# Perform OCR on detected objects
ocr_text= ""
for box in results[0].boxes.xyxy.cpu().numpy():
x1, y1, x2, y2 = map(int, box)
cropped = image[y1:y2, x1:x2]
# Skip if the cropped region is too small
if cropped.size == 0:
continue
# Apply detector for plate region
alpr = ALPR(detector_model=DETECTOR_MODELS[0], ocr_model=OCR_MODELS[0])
alpr_results = alpr.predict(cropped)
if alpr_results:
res = alpr_results[0]
# Access the detection and OCR attributes from ALPRResult
plate_text = res.ocr.text if res.ocr else "N/A"
plate_confidence = res.ocr.confidence if res.ocr else 0.0
ocr_text += f"- Detected Plate: {plate_text} with confidence {plate_confidence:.2f}\n"
# Add text to annotated image
cv2.putText(annotated_image, plate_text, (x1, y1-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
# Convert back to RGB for Gradio display
annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
return annotated_image, ocr_text
# Create Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# MIA-Yolov8 for peruvian plate recognition")
with gr.Row():
with gr.Column():
input_image = gr.Image(label="Input Image", type="pil")
conf_slider = gr.Slider(0, 1, value=0.25, label="Confidence Threshold")
iou_slider = gr.Slider(0, 1, value=0.45, label="IOU Threshold")
submit_btn = gr.Button("Run model")
with gr.Column():
output_image = gr.Image(label="Detected Objects")
ocr_output = gr.Textbox(label="OCR Results")
submit_btn.click(
fn=predict,
inputs=[input_image, conf_slider, iou_slider],
outputs=[output_image, ocr_output]
)
demo.launch() |