omarelsayeed's picture
Update app.py
30260da verified
from ultralytics import YOLO
import gradio as gr
from huggingface_hub import snapshot_download
from PIL import Image, ImageDraw, ImageFont
import cv2
import numpy as np
hoppa = {
"0": "٠",
"1": "١",
"2": "٢",
"3": "٣",
"4": "٤",
"5": "٥",
"6": "٦",
"7": "٧",
"8":"٨",
"9":"٩",
"Beh": "ب",
"Daad": "ض",
"Een": "ع",
"F": "ف",
"Heeh": "هه",
"Kaaf": "ك",
"Laam": "ل",
"License Plate": "",
"Meem": "م",
"Noon": "ن",
"Q": "ق",
"R": "ر",
"Saad": "ص",
"Seen": "س",
"Taa": "ط",
"Wow": "و",
"Yeeh": "ي",
"Zah": "ظ",
"Zeen": "ز",
"alef": "أ",
"car": "",
"daal": "د",
"geem": "ج",
"ghayn": "غ",
"khaa": "خ",
"sheen": "ش",
"teh": "ت",
"theh": "ث",
"zaal": "ذ",
"7aah" : "ح"
}
# Model path
model_dir = snapshot_download("omarelsayeed/licence_plates") + "/license_yolo8s_1024.pt"
# YOLO model initialization
model = YOLO(model_dir)
def predict_and_draw(img, conf_threshold, iou_threshold):
"""
Predict objects in an image, annotate with bounding boxes and markers,
and generate a concatenated class string.
"""
# Convert Gradio input to numpy array
image = np.array(img)
# Predict with YOLO
results = model.predict(image, conf=conf_threshold, iou=iou_threshold, verbose=False ,imgsz=1024,)
if not results:
return img, "No objects detected."
# Extract bounding boxes, classes, and sort them right to left
boxes_and_classes = []
for box, cls in zip(results[0].boxes.xyxy, results[0].boxes.cls):
x_min, y_min, x_max, y_max = map(int, box[:4])
class_name = model.names[int(cls)]
boxes_and_classes.append((x_max, y_min, class_name, (x_min, y_min, x_max, y_max)))
# Sort boxes by x_max (right-to-left order)
boxes_and_classes.sort(key=lambda b: b[0], reverse=True)
# Extract and process classes
processed_classes = []
for _, _, cls, _ in boxes_and_classes:
if cls.isalpha():
processed_classes.append(hoppa.get(cls, cls))
elif not cls.isdigit():
processed_classes.append(cls)
else:
processed_classes.append(cls)
# Join letters with '-' and digits separately
concatenated_classes = '-'.join([cls for cls in processed_classes if cls.isalpha()])
digits = ''.join([cls for cls in processed_classes if cls.isdigit()])
digits = ''.join([hoppa[str(c)] for c in digits])
digits = digits[::-1]
if digits:
concatenated_classes += f'-{digits}' if concatenated_classes else digits
# Annotate image
annotated_image = Image.fromarray(image)
draw = ImageDraw.Draw(annotated_image)
for _, _, _, (x_min, y_min, x_max, y_max) in boxes_and_classes:
# Draw bounding box
draw.rectangle([x_min, y_min, x_max, y_max], outline="green", width=2)
# Mark the top-right corner
draw.ellipse([x_max - 5, y_min - 5, x_max + 5, y_min + 5], fill="red")
return annotated_image, concatenated_classes
# Gradio interface
iface = gr.Interface(
fn=predict_and_draw,
inputs=[
gr.Image(type="pil", label="Upload Image"),
gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold")
],
outputs=[
gr.Image(type="pil", label="Annotated Image"),
gr.Textbox(label="Concatenated Classes")
],
title="License Plate Recognition.",
description="Upload an image to detect license plates."
)
# Launch the interface
iface.launch()