import gradio as gr
import easyocr
from transformers import pipeline
from PIL import Image
import numpy as np

# Initialize EasyOCR for Persian
reader = easyocr.Reader(['fa'])

# Initialize text processing pipeline
try:
    # استفاده از مدل محلی یا مدل‌های عمومی
    text_processor = pipeline("text-generation", model="gpt2")  # مدل جایگزین
except Exception as e:
    text_processor = None

def run_ocr(image):
    """استخراج متن از تصویر با EasyOCR"""
    try:
        if isinstance(image, Image.Image):
            image = np.array(image)
        results = reader.readtext(image)
        return " ".join([result[1] for result in results]) if results else "متنی یافت نشد!"
    except Exception as e:
        return f"خطا در OCR: {str(e)}"

def process_text(text):
    """پردازش ساده متن"""
    if text == "متنی یافت نشد!":
        return text
    
    # اگر پردازشگر متن وجود داشت از آن استفاده کن
    if text_processor:
        try:
            return text_processor(text, max_length=50)[0]['generated_text']
        except:
            return text  # اگر خطا رخ داد متن اصلی را برگردان
    return text  # اگر پردازشگر متن وجود نداشت

with gr.Blocks(title="سیستم OCR فارسی") as app:
    gr.Markdown("## استخراج متن فارسی از تصاویر")
    
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(label="تصویر ورودی", type="pil")
            btn = gr.Button("پردازش تصویر", variant="primary")
        
        with gr.Column():
            ocr_output = gr.Textbox(label="متن استخراج شده")
            processed_output = gr.Textbox(label="متن پردازش شده", visible=False)  # غیرفعال شده
    
    btn.click(
        fn=lambda x: (run_ocr(x), process_text(run_ocr(x))),
        inputs=img_input,
        outputs=[ocr_output, processed_output]
    )

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