app.py

import os
# این خط باید قبل از هرگونه import از torch یا ultralytics باشد
os.environ['PYTORCH_WEIGHTS_ONLY'] = 'false'

import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import random

# --- 1. Load the Model ---
# توجه: نام مدل به 'best.pt' برگردانده شد.
# مطمئن شوید که فایل 'best.pt' در کنار این اسکریپت وجود دارد.
try:
    model = YOLO("best.pt")
except Exception as e:
    print(f"Error loading model: {e}")
    # در صورت عدم وجود مدل، یک مدل پیش‌فرض یولو بارگذاری می‌شود تا برنامه اجرا شود
    model = YOLO("yolov8n.pt")


# --- 2. Define Class Names & Character Map ---
# برای این پروژه، لیست کلاس‌ها تمام کاراکترهای ممکن را پوشش می‌دهد.
class_names = [
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
    'Alef', 'Be', 'Te', 'Se', 'Jim', 'Dal', 'Sin', 'Shin',
    'Sad', 'Ta', 'Za', 'Eyn', 'Ghaf', 'Lam', 'Mim', 'Nun',
    'He', 'Vav', 'Pe', 'Zhe', 'plate', 'Ye', 'Ze'
]

# نقشه برای تبدیل ID کلاس به کاراکتر فارسی برای نمایش در خروجی متنی
char_map = {
    0: '۰', 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: 'ژ', 31: 'ی', 32: 'ز'
    # کلاس 30 (plate) از این نقشه حذف شده زیرا یک کاراکتر نیست.
}

# برای هر کلاس یک رنگ رندوم تعریف می‌کنیم.
class_colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(class_names))]


# --- 3. Main Image Processing Function ---
def detect_and_read_plate(input_image: Image.Image):
    """
    This function takes an input image, detects all objects (plate and characters),
    draws bounding boxes, sorts the characters, and returns the processed image
    along with the recognized license plate text.
    """
    if input_image is None:
        return None, ""

    # اجرای مدل روی تصویر ورودی
    results = model(input_image, verbose=False)[0]
    
    # تبدیل تصویر ورودی به فرمت مناسب برای پردازش با OpenCV
    img_np = np.array(input_image.convert("RGB"))
    img_copy = img_np.copy()

    recognized_characters = []

    # بررسی اینکه آیا باکسی تشخیص داده شده است یا خیر
    if results.boxes is not None:
        boxes = results.boxes.data.cpu().numpy()

        for box in boxes:
            x1, y1, x2, y2 = map(int, box[:4])
            confidence = box[4]
            cls_id = int(box[5])

            if cls_id < len(class_names):
                # گرفتن رنگ و لیبل برای رسم کادر
                color = class_colors[cls_id]
                label = f"{class_names[cls_id]} {confidence:.2f}"
                
                # رسم کادر دور هر شیء تشخیص داده شده
                cv2.rectangle(img_copy, (x1, y1), (x2, y2), color, 2)
                cv2.putText(img_copy, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

                # اگر شیء تشخیص داده شده یک کاراکتر بود، آن را برای مرتب‌سازی ذخیره کن
                if cls_id in char_map:
                    recognized_characters.append({'box': box, 'char': char_map[cls_id]})
    
    # مرتب‌سازی کاراکترهای تشخیص داده شده بر اساس موقعیت افقی (x1)
    plate_text = ""
    if recognized_characters:
        sorted_chars = sorted(recognized_characters, key=lambda item: item['box'][0])
        # ساختن رشته نهایی پلاک
        plate_text = " ".join([item['char'] for item in sorted_chars])
    else:
        plate_text = "No characters were detected."

    return Image.fromarray(img_copy), plate_text


# --- 4. Dynamically Read Examples from a Folder (unchanged) ---
# نام فولدر به نامی مرتبط با پروژه تغییر کرده است.
examples_folder = "Persian_License_Plate_Images"
examples_list = []
try:
    all_files_in_folder = os.listdir(examples_folder)
    # فیلتر کردن فایل‌ها برای اطمینان از اینکه فقط تصاویر خوانده می‌شوند
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
    image_files = [f for f in all_files_in_folder if os.path.splitext(f)[1].lower() in image_extensions]
    for filename in image_files:
        full_path = os.path.join(examples_folder, filename)
        examples_list.append(full_path)
except FileNotFoundError:
    print(f"Warning: Folder '{examples_folder}' not found. No examples will be shown.")
except Exception as e:
    print(f"An error occurred while reading examples: {e}")


# ===================================================================
# --- Key Section: Building the Custom Layout with gr.Blocks ---
# ===================================================================

# راهنمای متنی که برای پروژه تشخیص پلاک بازنویسی شده است.
guide_html = """
<hr>
<div style="text-align: center;">
  <h3>🚗 Iranian License Plate Detection & Reading Guide 🇮🇷</h3>
  <p>
    This model detects <strong>Iranian vehicle license plates</strong> and reads the characters on them.
    Upload an image, and the model will draw boxes around the plate and each character, then display the read text below.
  </p>
  <br>
  <p><strong>Note:</strong> For best results, use clear images where the license plate is easily readable.</p>
</div>
"""

# شروع تعریف بلوک‌های رابط کاربری
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    # ۱. عنوان اصلی و توضیحات اولیه
    gr.Markdown("<h1 style='text-align: center;'>Persian License Plate Detection and OCR</h1>")
    gr.Markdown(
        "**Welcome!** Upload a vehicle image or click an example. "
        "The model will draw bounding boxes and output the recognized plate text."
    )
    
    # ۲. چیدمان ورودی و خروجی
    with gr.Row():
        input_image = gr.Image(type="pil", label="Upload Vehicle Image")
        output_image = gr.Image(type="pil", label="Processed Image with Detections")
    
    # کادر متنی برای نمایش پلاک خوانده شده
    recognized_text = gr.Textbox(label="Recognized Plate Text", interactive=False)

    # ۳. دکمه برای ارسال تصویر
    submit_btn = gr.Button("🔍 Detect and Read Plate")

    # ۴. بخش نمونه‌ها
    if examples_list:
        gr.Examples(
            examples=examples_list,
            inputs=input_image,
            # خروجی‌ها شامل تصویر و کادر متنی هستند
            outputs=[output_image, recognized_text],
            fn=detect_and_read_plate,
            cache_examples=True,
        )

    # ۵. متن راهنما
    gr.Markdown(guide_html)
    
    # تعریف عملکرد دکمه: خروجی‌ها به تصویر و کادر متنی ارسال می‌شوند
    submit_btn.click(fn=detect_and_read_plate, inputs=input_image, outputs=[output_image, recognized_text])

# اجرای برنامه
if __name__ == "__main__":
    demo.launch(share=True)