Update app.py

app.py

@@ -1,48 +1,32 @@
 import gradio as gr
-from doctr.io import DocumentFile
 from doctr.models import ocr_predictor
-import fitz  # PyMuPDF
-import io
-from PIL import Image
-
-# Initialize the OCR model
-model = ocr_predictor(pretrained=True)
+from doctr.io import DocumentFile
 
-def perform_ocr(file):
-    if file.name.lower().endswith('.pdf'):
-        # Process PDF
-        text = ""
-        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:
-        # Process image
-        doc = DocumentFile.from_images(file.name)
-        result = model(doc)
-        return result.render()
+# Cargar el modelo OCR de DocTR
+ocr_model = ocr_predictor(pretrained=True)
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=perform_ocr,
-    inputs=gr.File(label="Upload PDF or Image"),
-    outputs="text",
-    title="OCR with doctr (PDF and Images)",
-    description="Upload a PDF file or an image to extract text using OCR."
-)
+# Función para procesar un PDF y extraer texto
+def extract_text_from_pdf(pdf_file):
+    # Leer el PDF con DocTR
+    doc = DocumentFile.from_pdf(pdf_file)
+    # Ejecutar el OCR
+    result = ocr_model(doc)
+    # Extraer el texto
+    text = "\n".join([block[1] for page in result.pages for block in page.blocks])
+    return text
 
-# Launch the interface
-iface.launch()
+# Crear la interfaz Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("# DocTR OCR para PDFs con Gradio")
+    gr.Markdown("Sube un archivo PDF para extraer el texto.")
+    
+    with gr.Row():
+        pdf_input = gr.File(label="Sube tu PDF", type="file")
+        text_output = gr.Textbox(label="Texto Extraído", lines=10)
+    
+    extract_button = gr.Button("Extraer Texto")
+    extract_button.click(extract_text_from_pdf, inputs=pdf_input, outputs=text_output)
 
+# Ejecutar la aplicación
+if __name__ == "__main__":
+    demo.launch()