import gradio as gr
from doctr.io import DocumentFile
from doctr.models import ocr_predictor

# Cargar el modelo preentrenado
model = ocr_predictor(pretrained=True)

def process_file(file):
    """Procesa un archivo (PDF o imagen) con docTR y retorna el texto extraído."""
    if file is None:
        return "Por favor, sube un archivo."
    
    # Leer el archivo subido
    doc = DocumentFile.from_pdf(file.name) if file.name.endswith('.pdf') else DocumentFile.from_images(file.name)

    # Realizar OCR
    result = model(doc)

    # Extraer el texto y retornarlo
    extracted_text = "\n".join([block['text'] for page in result.pages for block in page['blocks']])
    return extracted_text

# Configuración de la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("## OCR con docTR")
    gr.Markdown("Sube un archivo PDF o una imagen para extraer texto utilizando un modelo preentrenado de docTR.")

    with gr.Row():
        input_file = gr.File(label="Subir archivo (PDF o imagen)")
        output_text = gr.Textbox(label="Texto extraído", lines=10)

    process_button = gr.Button("Procesar archivo")

    process_button.click(fn=process_file, inputs=[input_file], outputs=[output_text])

# Ejecutar la app
if __name__ == "__main__":
    demo.launch()