Update app.py

app.py

@@ -1,52 +1,48 @@
 import gradio as gr
-import pytesseract
-from PIL import Image
-import os
-from pdf2image import convert_from_bytes
+from doctr.io import DocumentFile
+from doctr.models import ocr_predictor
+import fitz  # PyMuPDF
 import io
+from PIL import Image
 
-# Configurar Tesseract para usar el modelo entrenado en Hugging Face Spaces
-tessdata_dir = "/home/user/.apt/usr/share/tesseract-ocr/4.00/tessdata"
-if os.path.exists(tessdata_dir):
-    pytesseract.pytesseract.tesseract_cmd = '/home/user/.apt/usr/bin/tesseract'
-    os.environ["TESSDATA_PREFIX"] = tessdata_dir
+# Initialize the OCR model
+model = ocr_predictor(pretrained=True)
 
 def perform_ocr(file):
-    if file is None:
-        return "Por favor, sube un archivo."
-    
-    # Verificar si el archivo es un PDF
     if file.name.lower().endswith('.pdf'):
-        # Convertir PDF a imágenes
-        try:
-            images = convert_from_bytes(file.read() if hasattr(file, 'read') else file)
-        except Exception as e:
-            return f"Error al procesar el PDF: {str(e)}"
-        
+        # Process PDF
         text = ""
-        for image in images:
-            text += pytesseract.image_to_string(image) + "\n\n"
+        pdf_document = fitz.open(file.name)
+        for page_num in range(pdf_document.page_count):
+            page = pdf_document[page_num]
+            pix = page.get_pixmap()
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            
+            # Convert PIL Image to bytes
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format='PNG')
+            img_byte_arr = img_byte_arr.getvalue()
+            
+            # Perform OCR on the image
+            doc = DocumentFile.from_images(img_byte_arr)
+            result = model(doc)
+            text += result.render() + "\n\n"  # Add newlines between pages
+        return text.strip()
     else:
-        # Procesar como imagen
-        try:
-            if hasattr(file, 'read'):
-                image = Image.open(io.BytesIO(file.read()))
-            else:
-                image = Image.open(file)
-            text = pytesseract.image_to_string(image)
-        except Exception as e:
-            return f"Error al procesar la imagen: {str(e)}"
-    
-    return text
+        # Process image
+        doc = DocumentFile.from_images(file.name)
+        result = model(doc)
+        return result.render()
 
-# Crear la interfaz de Gradio
+# Create Gradio interface
 iface = gr.Interface(
     fn=perform_ocr,
-    inputs=gr.File(label="Sube una imagen o PDF"),
+    inputs=gr.File(label="Upload PDF or Image"),
     outputs="text",
-    title="Tesseract OCR para Imágenes y PDFs",
-    description="Sube una imagen o un archivo PDF para extraer texto usando Tesseract OCR."
+    title="OCR with doctr (PDF and Images)",
+    description="Upload a PDF file or an image to extract text using OCR."
 )
 
-# Lanzar la interfaz
+# Launch the interface
 iface.launch()
+