Update app.py

app.py

@@ -1,32 +1,43 @@
 import gradio as gr
-from doctr.models import ocr_predictor
+import io
 from doctr.io import DocumentFile
+from doctr.models import ocr_predictor
 
-# Cargar el modelo OCR de DocTR
-ocr_model = ocr_predictor(pretrained=True)
-
-# 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
+# Initialize the OCR model
+model = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True)
 
-# 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.")
+def ocr_process(file):
+    # Read the uploaded file
+    if file.name.lower().endswith('.pdf'):
+        doc = DocumentFile.from_pdf(file.name)
+    else:
+        # Assume it's an image if not PDF
+        image_stream = io.BytesIO(file.read())
+        doc = DocumentFile.from_images(image_stream)
     
-    with gr.Row():
-        pdf_input = gr.File(label="Sube tu PDF", type="file")
-        text_output = gr.Textbox(label="Texto Extraído", lines=10)
+    # Perform OCR
+    result = model(doc)
     
-    extract_button = gr.Button("Extraer Texto")
-    extract_button.click(extract_text_from_pdf, inputs=pdf_input, outputs=text_output)
+    # Extract text from the result
+    extracted_text = ""
+    for page in result.pages:
+        for block in page.blocks:
+            for line in block.lines:
+                for word in line.words:
+                    extracted_text += word.value + " "
+                extracted_text += "\n"
+            extracted_text += "\n"
+    
+    return extracted_text.strip()
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=ocr_process,
+    inputs=gr.File(label="Upload PDF or Image"),
+    outputs=gr.Textbox(label="Extracted Text"),
+    title="OCR with doctr",
+    description="Upload a PDF or image file to extract text using OCR."
+)
 
-# Ejecutar la aplicación
-if __name__ == "__main__":
-    demo.launch()
+# Launch the interface
+iface.launch()