app.py

import gradio as gr
import time
import numpy as np
from PIL import Image
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import easyocr
from doctr.models import ocr_predictor

# Initialize models
models = {
    "EasyOCR": easyocr.Reader(['fa']),  # تنظیم زبان فارسی
    "TrOCR": {
        "processor": TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed"),
        "model": VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-printed")
    },
    "DocTR": ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True)
}

def run_easyocr(image):
    """تابع پردازش تصویر با EasyOCR برای فارسی"""
    try:
        result = models["EasyOCR"].readtext(np.array(image), detail=0)
        return ' '.join(result) if result else ''
    except Exception as e:
        return f"خطا: {str(e)}"

def run_trocr(image):
    """تابع پردازش تصویر با TrOCR برای فارسی"""
    try:
        pixel_values = models["TrOCR"]["processor"](image, return_tensors="pt").pixel_values
        generated_ids = models["TrOCR"]["model"].generate(pixel_values)
        return models["TrOCR"]["processor"].batch_decode(generated_ids, skip_special_tokens=True)[0]
    except Exception as e:
        return f"خطا: {str(e)}"

def run_doctr(image):
    """تابع پردازش تصویر با DocTR برای فارسی"""
    try:
        if isinstance(image, Image.Image):
            image = np.array(image)
        result = models["DocTR"]([image])
        return ' '.join([word[0] for page in result.pages for block in page.blocks 
                        for line in block.lines for word in line.words])
    except Exception as e:
        return f"خطا: {str(e)}"

def compare_models(image):
    """تابع اصلی مقایسه مدل‌ها"""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.convert("RGB")
    
    results = {}
    times = {}
    
    # اجرای تمام مدل‌های OCR
    for name, func in [("EasyOCR", run_easyocr),
                      ("TrOCR", run_trocr),
                      ("DocTR", run_doctr)]:
        start = time.time()
        results[name] = func(image)
        times[name] = time.time() - start
    
    # ایجاد جدول مقایسه
    table_rows = []
    for name in results:
        table_rows.append(f"""
        <tr>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: center; font-weight: bold;">{name}</td>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: right; direction: rtl;">{results[name]}</td>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: center;">{times[name]:.3f} ثانیه</td>
        </tr>
        """)
    
    comparison = f"""
    <div style="overflow-x: auto;">
    <table style="width:100%; border-collapse: collapse; margin: 15px 0; font-family: Arial, sans-serif;">
        <tr style="background-color: #4CAF50; color: white;">
            <th style="padding: 12px; border: 1px solid #ddd; text-align: center;">مدل</th>
            <th style="padding: 12px; border: 1px solid #ddd; text-align: center;">متن استخراج شده</th>
            <th style="padding: 12px; border: 1px solid #ddd; text-align: center;">زمان پردازش</th>
        </tr>
        {''.join(table_rows)}
    </table>
    </div>
    """
    
    return comparison, results['EasyOCR'], results['TrOCR'], results['DocTR']

# رابط کاربری Gradio
with gr.Blocks(title="مقایسه مدل‌های OCR فارسی", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🚀 مقایسه مدل‌های تشخیص متن فارسی
    مقایسه عملکرد مدل‌های مختلف OCR برای استخراج متن از تصاویر فارسی
    """)
    
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(label="تصویر ورودی", type="pil")
            gr.Examples(
                examples=["sample_fa1.jpg", "sample_fa2.png"],
                inputs=img_input,
                label="تصاویر نمونه"
            )
            submit_btn = gr.Button("مقایسه مدل‌ها", variant="primary")
        
        with gr.Column():
            comparison = gr.HTML(label="نتایج مقایسه")
            with gr.Accordion("نتایج تفکیکی", open=False):
                gr.Markdown("### خروجی هر مدل")
                easy_output = gr.Textbox(label="EasyOCR")
                trocr_output = gr.Textbox(label="TrOCR")
                doctr_output = gr.Textbox(label="DocTR")
    
    submit_btn.click(
        fn=compare_models,
        inputs=img_input,
        outputs=[comparison, easy_output, trocr_output, doctr_output]
    )

if __name__ == "__main__":
    demo.launch()