Update app.py

app.py

@@ -1,57 +1,112 @@
 import gradio as gr
-from paddleocr import PaddleOCR
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import time
 from PIL import Image
-import os
+from paddleocr import PaddleOCR
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import pytesseract
+import numpy as np
 
-# --- مدل‌ها ---
-try:
-    model_name = "m3hrdadfi/mt5-small-finetuned-grammar-synthesis"
-    
-    ocr_model = PaddleOCR(lang='fa', use_textline_orientation=True)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    nlp_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    
-except Exception as e:
-    raise gr.Error(f"خطا در بارگذاری مدل‌ها: لطفاً این خطا را به توسعه دهنده گزارش دهید. خطا: {str(e)}")
+# Initialize models
+paddle_ocr = PaddleOCR(lang='fa', use_textline_orientation=True)
+trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
+trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-printed")
 
-# --- توابع پردازش ---
-def run_ocr(image):
-    image_path = image.name
-    result = ocr_model.ocr(image_path, cls=True)
-    texts = [line[1][0] for line in result[0]] if result else []
-    os.remove(image_path)
-    return " ".join(texts)
+def run_paddleocr(image):
+    """Run PaddleOCR on image"""
+    image_path = "temp.jpg"
+    image.save(image_path)
+    result = paddle_ocr.ocr(image_path, cls=True)
+    text = ' '.join([line[1][0] for line in result[0]]) if result else ''
+    return text
 
-def postprocess_text(text):
-    inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
-    outputs = nlp_model.generate(**inputs)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def run_trocr(image):
+    """Run TrOCR on image"""
+    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]
+
+def run_tesseract(image):
+    """Run Tesseract OCR on image"""
+    return pytesseract.image_to_string(image, lang='fas')
+
+def compare_models(image):
+    """Compare all three OCR models"""
+    # Convert to RGB if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    image = image.convert("RGB")
+    
+    results = {}
+    
+    # Run PaddleOCR
+    start = time.time()
+    results['PaddleOCR'] = run_paddleocr(image)
+    paddle_time = time.time() - start
+    
+    # Run TrOCR
+    start = time.time()
+    results['TrOCR'] = run_trocr(image)
+    trocr_time = time.time() - start
+    
+    # Run Tesseract
+    start = time.time()
+    results['Tesseract'] = run_tesseract(image)
+    tesseract_time = time.time() - start
+    
+    # Create comparison table
+    comparison = f"""
+    <table>
+        <tr>
+            <th>مدل</th>
+            <th>متن استخراج شده</th>
+            <th>زمان پردازش (ثانیه)</th>
+        </tr>
+        <tr>
+            <td>PaddleOCR</td>
+            <td>{results['PaddleOCR']}</td>
+            <td>{paddle_time:.2f}</td>
+        </tr>
+        <tr>
+            <td>TrOCR</td>
+            <td>{results['TrOCR']}</td>
+            <td>{trocr_time:.2f}</td>
+        </tr>
+        <tr>
+            <td>Tesseract</td>
+            <td>{results['Tesseract']}</td>
+            <td>{tesseract_time:.2f}</td>
+        </tr>
+    </table>
+    """
+    
+    return comparison, results['PaddleOCR'], results['TrOCR'], results['Tesseract']
 
-# --- رابط کاربری ---
-with gr.Blocks() as app:
-    gr.Markdown("## سیستم OCR فارسی با پردازش پیشرفته متن")
+# Create Gradio interface
+with gr.Blocks(title="مقایسه مدل‌های OCR فارسی") as demo:
+    gr.Markdown("""
+    ## مقایسه عملکرد مدل‌های OCR برای زبان فارسی
+    این برنامه سه مدل مختلف OCR را روی تصاویر فارسی مقایسه می‌کند:
+    1. PaddleOCR
+    2. TrOCR (مایکروسافت)
+    3. Tesseract OCR
+    """)
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="filepath", label="تصویر ورودی")
-            process_btn = gr.Button("پردازش تصویر")
+            image_input = gr.Image(label="تصویر ورودی", type="pil")
+            submit_btn = gr.Button("مقایسه مدل‌ها")
         
         with gr.Column():
-            raw_output = gr.Textbox(label="متن استخراج شده")
-            processed_output = gr.Textbox(label="متن پردازش شده")
-    
-    def process_image(img):
-        # Avoid running OCR twice
-        raw = run_ocr(img)
-        processed = postprocess_text(raw)
-        return raw, processed
-
-    process_btn.click(
-        fn=process_image,
+            comparison_output = gr.HTML(label="نتایج مقایسه")
+            paddle_output = gr.Textbox(label="PaddleOCR")
+            trocr_output = gr.Textbox(label="TrOCR")
+            tesseract_output = gr.Textbox(label="Tesseract")
+    
+    submit_btn.click(
+        fn=compare_models,
         inputs=image_input,
-        outputs=[raw_output, processed_output]
+        outputs=[comparison_output, paddle_output, trocr_output, tesseract_output]
     )
 
 if __name__ == "__main__":
-    app.launch()
+    demo.launch()