import gradio as gr
import time
import numpy as np
from PIL import Image
from paddleocr import PaddleOCR
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import easyocr

# Initialize models
paddle_ocr = PaddleOCR(lang='en')  # PaddleOCR برای انگلیسی
easy_ocr = easyocr.Reader(['en'])  # EasyOCR برای انگلیسی
trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-printed")

def run_paddleocr(image):
    """Run PaddleOCR on image"""
    if isinstance(image, Image.Image):
        image = np.array(image)
    
    try:
        result = paddle_ocr.ocr(image)
        return ' '.join([line[1][0] for line in result[0]]) if result else ''
    except Exception as e:
        return f"PaddleOCR Error: {str(e)}"

def run_easyocr(image):
    """Run EasyOCR on image"""
    if isinstance(image, Image.Image):
        image = np.array(image)
    
    try:
        result = easy_ocr.readtext(image, detail=0)
        return ' '.join(result) if result else ''
    except Exception as e:
        return f"EasyOCR Error: {str(e)}"

def run_trocr(image):
    """Run TrOCR on image"""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    try:
        pixel_values = trocr_processor(image, return_tensors="pt").pixel_values
        generated_ids = trocr_model.generate(pixel_values)
        return trocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    except Exception as e:
        return f"TrOCR Error: {str(e)}"

def compare_models(image):
    """Compare all three OCR models"""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.convert("RGB")
    
    results = {}
    times = {}
    
    # Run PaddleOCR
    start = time.time()
    results['PaddleOCR'] = run_paddleocr(image)
    times['PaddleOCR'] = time.time() - start
    
    # Run EasyOCR
    start = time.time()
    results['EasyOCR'] = run_easyocr(image)
    times['EasyOCR'] = time.time() - start
    
    # Run TrOCR
    start = time.time()
    results['TrOCR'] = run_trocr(image)
    times['TrOCR'] = time.time() - start
    
    # Create comparison table
    comparison = f"""
    <table style="width:100%; border-collapse: collapse;">
        <tr style="background-color: #f2f2f2;">
            <th style="padding: 8px; border: 1px solid #ddd; text-align: center;">Model</th>
            <th style="padding: 8px; border: 1px solid #ddd; text-align: center;">Extracted Text</th>
            <th style="padding: 8px; border: 1px solid #ddd; text-align: center;">Processing Time (s)</th>
        </tr>
        <tr>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: center;">PaddleOCR</td>
            <td style="padding: 8px; border: 1px solid #ddd;">{results['PaddleOCR']}</td>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: center;">{times['PaddleOCR']:.3f}</td>
        </tr>
        <tr>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: center;">EasyOCR</td>
            <td style="padding: 8px; border: 1px solid #ddd;">{results['EasyOCR']}</td>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: center;">{times['EasyOCR']:.3f}</td>
        </tr>
        <tr>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: center;">TrOCR</td>
            <td style="padding: 8px; border: 1px solid #ddd;">{results['TrOCR']}</td>
            <td style="padding: 8px; border: 1px solid #ddd; text-align: center;">{times['TrOCR']:.3f}</td>
        </tr>
    </table>
    """
    
    return comparison, results['PaddleOCR'], results['EasyOCR'], results['TrOCR']

# Create Gradio interface
with gr.Blocks(title="English OCR Comparison Tool") as demo:
    gr.Markdown("""
    ## English OCR Models Comparison
    This tool compares three OCR models for English text:
    1. PaddleOCR
    2. EasyOCR
    3. TrOCR (Microsoft)
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Input Image", type="pil")
            submit_btn = gr.Button("Compare Models", variant="primary")
        
        with gr.Column():
            comparison_output = gr.HTML(label="Comparison Results")
            with gr.Accordion("Individual Results", open=False):
                paddle_output = gr.Textbox(label="PaddleOCR Result")
                easy_output = gr.Textbox(label="EasyOCR Result")
                trocr_output = gr.Textbox(label="TrOCR Result")
    
    submit_btn.click(
        fn=compare_models,
        inputs=image_input,
        outputs=[comparison_output, paddle_output, easy_output, trocr_output]
    )

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