Update app.py

app.py

@@ -1,122 +1,90 @@
 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
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from PIL import Image
+import numpy as np
+import os
 
-# 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)
-}
+# Initialize EasyOCR for Persian
+reader = easyocr.Reader(['fa'])
 
-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)}"
+# 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_trocr(image):
-    """تابع پردازش تصویر با TrOCR برای فارسی"""
+def run_ocr(image):
+    """استخراج متن از تصویر با EasyOCR"""
     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]
+        # 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)}"
+        return f"خطا در پردازش تصویر: {str(e)}"
 
-def run_doctr(image):
-    """تابع پردازش تصویر با DocTR برای فارسی"""
+def postprocess_text(text):
+    """پردازش متن با مدل زبانی"""
     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])
+        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)}"
+        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
+def process_image(image):
+    """پردازش کامل تصویر"""
+    # استخراج متن
+    raw_text = run_ocr(image)
     
-    # ایجاد جدول مقایسه
-    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>
-        """)
+    # پردازش متن
+    processed_text = postprocess_text(raw_text) if raw_text and raw_text != "متنی یافت نشد!" else raw_text
     
-    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']
+    return raw_text, processed_text
 
 # رابط کاربری Gradio
-with gr.Blocks(title="مقایسه مدل‌های OCR فارسی", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="OCR فارسی با پردازش NLP") as app:
     gr.Markdown("""
-    # 🚀 مقایسه مدل‌های تشخیص متن فارسی
-    مقایسه عملکرد مدل‌های مختلف OCR برای استخراج متن از تصاویر فارسی
+    ## 🔠 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")
+            image_input = gr.Image(label="تصویر حاوی متن فارسی را آپلود کنید", type="pil")
+            process_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")
+            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]
+    )
     
-    submit_btn.click(
-        fn=compare_models,
-        inputs=img_input,
-        outputs=[comparison, easy_output, trocr_output, doctr_output]
+    # نمایش خودکار تصویر ورودی
+    image_input.change(
+        fn=lambda img: img,
+        inputs=image_input,
+        outputs=image_output
     )
+    
+    gr.Markdown("---")
+    gr.Markdown("""
+    **راهنما**:
+    1. تصویری حاوی متن فارسی آپلود کنید
+    2. روی دکمه 'پردازش تصویر' کلیک کنید
+    3. نتایج استخراج متن و پردازش زبان طبیعی را مشاهده کنید
+    """)
 
 if __name__ == "__main__":
-    demo.launch()
+    app.launch()