Update app.py

app.py

@@ -7,33 +7,51 @@ from ultralytics import YOLO
 import cv2
 import numpy as np
 from PIL import Image
+import random
 
 # --- 1. Load the Model ---
-# توجه: نام مدل به چیزی که برای تشخیص پلاک آموزش داده‌اید تغییر کرده است.
-# مطمئن شوید که فایل 'persian_license_plate_v1.pt' در کنار این اسکریپت وجود دارد.
+# توجه: نام مدل به 'best.pt' برگردانده شد.
+# مطمئن شوید که فایل 'best.pt' در کنار این اسکریپت وجود دارد.
 try:
-    model = YOLO("persian_license_plate_v1.pt")
+    model = YOLO("best.pt")
 except Exception as e:
     print(f"Error loading model: {e}")
     # در صورت عدم وجود مدل، یک مدل پیش‌فرض یولو بارگذاری می‌شود تا برنامه اجرا شود
     model = YOLO("yolov8n.pt")
 
 
-# --- 2. Define Class Names ---
-# برای این پروژه، ما فقط یک کلاس "پلاک خودرو" داریم.
-class_names = ['license_plate']
-# برای هر کلاس یک رنگ تعریف می‌کنیم (در اینجا فقط یک رنگ لازم است).
-class_colors = [(0, 255, 0)] # رنگ سبز برای پلاک
+# --- 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'
+]
 
-# --- 3. Main Image Processing Function (Core logic unchanged) ---
-def detect_license_plate(input_image: Image.Image):
+# نقشه برای تبدیل 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 a license plate,
-    draws a bounding box around it, and returns the result.
+    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
-    
+        return None, ""
+
     # اجرای مدل روی تصویر ورودی
     results = model(input_image, verbose=False)[0]
     
@@ -41,30 +59,43 @@ def detect_license_plate(input_image: Image.Image):
     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):
-                cls_name = class_names[cls_id]
+                # گرفتن رنگ و لیبل برای رسم کادر
                 color = class_colors[cls_id]
+                label = f"{class_names[cls_id]} {confidence:.2f}"
                 
-                # استخراج مختصات کادر
-                x1, y1, x2, y2 = map(int, box[:4])
-                
-                # رسم کادر دور پلاک
-                cv2.rectangle(img_copy, (x1, y1), (x2, y2), color, 3)
-                
-                # آماده‌سازی و نوشتن لیبل (نام کلاس و امتیاز اطمینان)
-                label = f"{cls_name} {box[4]:.2f}"
-                cv2.putText(img_copy, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
+                # رسم کادر دور هر شیء تشخیص داده شده
+                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)
+    return Image.fromarray(img_copy), plate_text
 
-# --- 4. Dynamically Read Examples from a Folder ---
+
+# --- 4. Dynamically Read Examples from a Folder (unchanged) ---
 # نام فولدر به نامی مرتبط با پروژه تغییر کرده است.
 examples_folder = "Persian_License_Plate_Images"
 examples_list = []
@@ -81,6 +112,7 @@ except FileNotFoundError:
 except Exception as e:
     print(f"An error occurred while reading examples: {e}")
 
+
 # ===================================================================
 # --- Key Section: Building the Custom Layout with gr.Blocks ---
 # ===================================================================
@@ -89,56 +121,52 @@ except Exception as e:
 guide_html = """
 <hr>
 <div style="text-align: center;">
-  <h3>🚗 Iranian License Plate Detection Guide 🇮🇷</h3>
+  <h3>🚗 Iranian License Plate Detection & Reading Guide 🇮🇷</h3>
   <p>
-    This model is trained to detect and draw a box around <strong>standard Iranian vehicle license plates</strong>.
-    Upload your image, and the model will automatically identify the plate's location.
+    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>
-  <div style="display: flex; justify-content: center; align-items: center; gap: 20px; margin-top: 15px;">
-    <div style="padding: 10px; border: 1px solid #4CAF50; border-radius: 8px; background-color: #f2fff2;">
-      <strong style="color: #005704;">🟩 License Plate</strong><br>
-      <em>
-        The green bounding box indicates the detected location of the license plate, along with the model's confidence score.
-      </em>
-    </div>
-  </div>
   <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="green")) as demo:
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     # ۱. عنوان اصلی و توضیحات اولیه
-    gr.Markdown("<h1 style='text-align: center;'>Persian License Plate Detection</h1>")
+    gr.Markdown("<h1 style='text-align: center;'>Persian License Plate Detection and OCR</h1>")
     gr.Markdown(
-        "**Welcome!** Upload an image of a vehicle or click on one of the examples below. "
-        "The model will draw a green bounding box around the license plate and show the predicted class (`license_plate`)."
+        "**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 Detected Plate")
+        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 License Plate")
+    submit_btn = gr.Button("🔍 Detect and Read Plate")
 
     # ۴. بخش نمونه‌ها
     if examples_list:
         gr.Examples(
             examples=examples_list,
             inputs=input_image,
-            outputs=output_image,
-            fn=detect_license_plate,
+            # خروجی‌ها شامل تصویر و کادر متنی هستند
+            outputs=[output_image, recognized_text],
+            fn=detect_and_read_plate,
             cache_examples=True,
         )
 
-    # ۵. متن راهنمای کلاس‌ها که در پایین صفحه قرار می‌گیرد
+    # ۵. متن راهنما
     gr.Markdown(guide_html)
     
-    # تعریف عملکرد دکمه: وقتی کلیک شد، تابع پردازش را با ورودی و خروجی مناسب اجرا کن
-    submit_btn.click(fn=detect_license_plate, inputs=input_image, outputs=output_image)
+    # تعریف عملکرد دکمه: خروجی‌ها به تصویر و کادر متنی ارسال می‌شوند
+    submit_btn.click(fn=detect_and_read_plate, inputs=input_image, outputs=[output_image, recognized_text])
 
 # اجرای برنامه
 if __name__ == "__main__":