import gradio as gr
import easyocr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from PIL import Image
import numpy as np
import os

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

# Load NLP model for text correction
tokenizer = AutoTokenizer.from_pretrained("persiannlp/mt5-small-parsinlu-grammar-correction")
model = AutoModelForSeq2SeqLM.from_pretrained("persiannlp/mt5-small-parsinlu-grammar-correction")

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

def postprocess_text(text):
    """پردازش متن با مدل زبانی"""
    try:
        inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
        outputs = model.generate(**inputs)
        return tokenizer.decode(outputs[0], skip_special_tokens=True)
    except Exception as e:
        return f"خطا در پردازش متن: {str(e)}"

def process_image(image):
    """پردازش کامل تصویر"""
    # استخراج متن
    raw_text = run_ocr(image)
    
    # پردازش متن
    processed_text = postprocess_text(raw_text) if raw_text and raw_text != "متنی یافت نشد!" else raw_text
    
    return raw_text, processed_text

# رابط کاربری Gradio
with gr.Blocks(title="OCR فارسی با پردازش NLP") as app:
    gr.Markdown("""
    ## 🔠 OCR فارسی + پردازش متن با مدل زبانی
    متن را از تصاویر استخراج کنید و با مدل زبانی اصلاح کنید!
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="تصویر حاوی متن فارسی را آپلود کنید", type="pil")
            process_btn = gr.Button("پردازش تصویر", variant="primary")
        
        with gr.Column():
            with gr.Tab("نتایج"):
                raw_output = gr.Textbox(label="متن خام (OCR)")
                processed_output = gr.Textbox(label="متن پردازش‌شده (NLP)")
            
            with gr.Tab("تصویر"):
                image_output = gr.Image(label="تصویر ورودی")

    # پردازش هنگام کلیک یا آپلود تصویر
    process_btn.click(
        fn=process_image,
        inputs=image_input,
        outputs=[raw_output, processed_output]
    )
    
    # نمایش خودکار تصویر ورودی
    image_input.change(
        fn=lambda img: img,
        inputs=image_input,
        outputs=image_output
    )
    
    gr.Markdown("---")
    gr.Markdown("""
    **راهنما**:
    1. تصویری حاوی متن فارسی آپلود کنید
    2. روی دکمه 'پردازش تصویر' کلیک کنید
    3. نتایج استخراج متن و پردازش زبان طبیعی را مشاهده کنید
    """)

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