app.py

import streamlit as st
import numpy as np
import cv2
import os
import uuid
import time
from plate_detection import load_models, preprocess_image, detect_license_plate, crop_license_plate, highlight_license_plate, encode_image, batch_process_plates
from crnn_recognition import load_char_to_id_mapping, preprocess_plate, predict_and_decode_plate
from PIL import Image, ImageFont, ImageDraw
import io

# صفحه را پیکربندی و سبک RTL و فونت فارسی را اعمال می‌کنیم
st.set_page_config(page_title="سامانه تشخیص پلاک", layout="wide")
st.markdown(
    """
    <style>
    @import url('https://fonts.googleapis.com/css2?family=Vazirmatn:wght@300;400;600&display=swap');

    html, body, [class*="css"], .stApp {
      font-family: 'Vazirmatn', sans-serif !important;
      direction: rtl;
      text-align: right;
    }

    /* هم‌راستاسازی ورودی‌ها */
    textarea, input, .stTextInput input, .stNumberInput input, .stSelectbox div, .stTextArea textarea {
      direction: rtl !important;
      text-align: right !important;
    }

    /* دکمه‌ها */
    .stButton > button {
      font-family: 'Vazirmatn', sans-serif !important;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# Create directories if they don't exist
UPLOAD_FOLDER = "uploads"
CROPPED_FOLDER = "cropped"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(CROPPED_FOLDER, exist_ok=True)

# Load models (do this only once)
@st.cache_resource
def load_models_cached():
    yolo_model, crnn_model = load_models('models/yolo11s_28feb.pt', 'models/crnn_model_pred_20250225_102044.keras')
    char_to_id_mapping = load_char_to_id_mapping('models/char_to_id_mapping.csv')
    return yolo_model, crnn_model, char_to_id_mapping

yolo_model, crnn_model, char_to_id_mapping = load_models_cached()

def process_image(image_path, max_plates=3, conf_threshold=0.5):
    # Start timing
    start_time = time.time()
    
    # Load the image from the file path
    image = cv2.imread(image_path)

    if image is None:
        st.error("خطا در بارگذاری تصویر. لطفاً از معتبر بودن فایل اطمینان حاصل کنید.")
        return None, None, None

    # Preprocess the image
    image = preprocess_image(image)

    # Detect license plates with speed settings
    all_plates = detect_license_plate(image, yolo_model, conf_threshold=conf_threshold, max_plates=max_plates)
    
    if not all_plates:
        st.warning("هیچ پلاکی در تصویر شناسایی نشد.")
        return [], None, None, []

    # Crop all plates first
    cropped_plates = []
    plate_results = []
    cropped_images = []
    
    for i, plate_info in enumerate(all_plates):
        box = plate_info['box']
        confidence = plate_info['confidence']
        
        # Crop the license plate
        cropped_plate = crop_license_plate(image, box)
        
        if cropped_plate is not None:
            # Save the cropped plate
            cropped_filename = os.path.join(CROPPED_FOLDER, f"cropped_{uuid.uuid4()}.jpg")
            cv2.imwrite(cropped_filename, cropped_plate)
            cropped_images.append(cropped_filename)
            cropped_plates.append(cropped_plate)
            
            plate_results.append({
                'confidence': confidence,
                'box': box,
                'cropped_image': cropped_filename
            })
        else:
            plate_results.append({
                'confidence': confidence,
                'box': box,
                'cropped_image': None
            })
    
    # Batch process all plates for text recognition (much faster)
    if cropped_plates:
        try:
            text_results = batch_process_plates(cropped_plates, crnn_model, char_to_id_mapping)
            
            # Combine results
            for i, result in enumerate(plate_results):
                if i < len(text_results):
                    result['plate_number'] = text_results[i]['predicted_seq']
                else:
                    result['plate_number'] = 'قابل خواندن نیست'
        except Exception as e:
            st.error(f"خطا در بازشناسی متن: {str(e)}")
            # Fallback to individual processing
            for result in plate_results:
                result['plate_number'] = 'خطا در بازشناسی'
    else:
        for result in plate_results:
            result['plate_number'] = 'برش امکان‌پذیر نیست'

    # Highlight all license plates
    highlighted_image = highlight_license_plate(image, all_plates)
    
    # Encode the highlighted image
    highlighted_image_base64 = encode_image(highlighted_image)
    
    # Calculate processing time
    processing_time = time.time() - start_time
    
    return plate_results, highlighted_image_base64, highlighted_image, cropped_images, processing_time

def display_predicted_plate(plate_number):
    """Displays the predicted plate number with each character in a separate box."""
    if not plate_number:
        st.write("هیچ پلاکی شناسایی نشد.")
        return

    num_chars = len(plate_number)

    # Load Vazirmatn font
    st.markdown(
        """
        <style>
        @font-face {
            font-family: 'Vazirmatn';
            src: url('https://cdn.jsdelivr.net/gh/rastikerdar/vazirmatn@29.1.0/dist/Vazirmatn.woff2') format('woff2');
        }
        </style>
        """,
        unsafe_allow_html=True,
    )

    # Use columns to create the boxes
    cols = st.columns(num_chars)  # Create n columns
    # In RTL pages, columns render right-to-left. Reverse mapping to preserve original visual order.
    visual_cols = cols[::-1]

    for i, char in enumerate(plate_number):
        with visual_cols[i]:
            st.markdown(
                f"<div style='border: 2px solid black; "
                f"text-align: center; "
                f"background-color: #007bff; "
                f"font-family: Vazirmatn; "
                f"font-size: 2em; "
                f"padding: 5px; "
                f"width: 100%; "
                f"box-sizing: border-box; "
                f"direction: ltr; unicode-bidi: plaintext;'>"
                f"{char}</div>",
                unsafe_allow_html=True,
            )

def display_predicted_plates(plate_results):
    """Displays all predicted plate numbers with confidence scores."""
    if not plate_results:
        st.write("هیچ پلاکی شناسایی نشد.")
        return

    st.subheader(f"{len(plate_results)} پلاک شناسایی شد")
    
    for i, result in enumerate(plate_results):
        st.write(f"**پلاک {i+1}:** {result['plate_number']} (اطمینان: {result['confidence']:.2f})")
        
        # Display cropped image if available
        if result['cropped_image']:
            st.image(result['cropped_image'], caption=f"برش پلاک {i+1}", use_container_width=True)
        
        # Display individual characters in boxes
        if result['plate_number'] and result['plate_number'] != "قابل خواندن نیست":
            display_predicted_plate(result['plate_number'])

# Streamlit UI
st.title("سامانه تشخیص پلاک")

# Speed configuration sidebar
with st.sidebar:
    st.header("⚡ تنظیمات سرعت")
    speed_mode = st.selectbox(
        "حالت پردازش",
        ["متعادل", "سریع", "دقیق"],
        help="سریع: آستانه اطمینان کمتر، پلاک‌های کمتر. دقیق: آستانه بالاتر، پلاک‌های بیشتر."
    )
    
    if speed_mode == "سریع":
        max_plates = 2
        conf_threshold = 0.4
    elif speed_mode == "دقیق":
        max_plates = 5
        conf_threshold = 0.7
    else:  # متعادل
        max_plates = 3
        conf_threshold = 0.5
    
    st.info(f"حالت: {speed_mode}\nحداکثر پلاک: {max_plates}\nآستانه اطمینان: {conf_threshold}")

uploaded_file = st.file_uploader("بارگذاری تصویر", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Save the uploaded image
    image = Image.open(uploaded_file)
    if image.mode == 'RGBA':
        image = image.convert('RGB')
    image_filename = os.path.join(UPLOAD_FOLDER, f"uploaded_{uuid.uuid4()}.jpg")
    image.save(image_filename)

    # Show processing progress
    with st.spinner("در حال پردازش تصویر..."):
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        # Process the image with speed settings
        plate_results, highlighted_image_base64, highlighted_image, cropped_images, processing_time = process_image(
            image_filename, max_plates=max_plates, conf_threshold=conf_threshold
        )
        
        progress_bar.progress(100)
        status_text.text("✅ پردازش کامل شد!")

    # Show processing time
    st.success(f"⚡ پردازش در {processing_time:.2f} ثانیه انجام شد")

    # Display all predicted plates
    display_predicted_plates(plate_results)

    # Display the highlighted image
    st.subheader("پلاک‌های شناسایی‌شده")
    st.image(highlighted_image, use_container_width=True, channels="BGR")