DocQA_Agent / app.py
OmidSakaki's picture
Update app.py
5c3f634 verified
raw
history blame
4.8 kB
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()