pdf-to-markdown / app.py
Biifruu's picture
Update app.py
3e3d3c7 verified
raw
history blame
2.13 kB
import spaces
import gradio as gr
import fitz # PyMuPDF
import os
import tempfile
import ocrmypdf
def extract_text_markdown(doc):
markdown_output = ""
image_dir = "extracted_images"
os.makedirs(image_dir, exist_ok=True)
image_counter = 0
for page in doc:
blocks = page.get_text("dict")["blocks"]
elements = []
for b in blocks:
y = b["bbox"][1]
if b["type"] == 0: # Texto
for line in b["lines"]:
line_y = line["bbox"][1]
line_text = " ".join([span["text"] for span in line["spans"]]).strip()
if line_text:
elements.append((line_y, line_text))
elif b["type"] == 1: # Imagen
elements.append((y, "![imagen]()"))
# Ordenar elementos por posición vertical
elements.sort(key=lambda x: x[0])
# Reconstruir markdown respetando el espaciado visual
previous_y = None
for y, content in elements:
if previous_y is not None and abs(y - previous_y) > 10: # si hay espacio entre líneas, añadir salto
markdown_output += "\n"
markdown_output += content + "\n"
previous_y = y
markdown_output += "\n---\n\n" # separador entre páginas
return markdown_output.strip()
@spaces.GPU
def convert(pdf_file):
original_doc = fitz.open(pdf_file)
plain_text = "\n".join([page.get_text() for page in original_doc])
# Si es imagen escaneada sin texto, aplicamos OCR
if len(plain_text.strip()) < 100:
ocr_temp_path = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False).name
ocrmypdf.ocr(pdf_file, ocr_temp_path, force_ocr=True)
doc = fitz.open(ocr_temp_path)
else:
doc = original_doc
markdown = extract_text_markdown(doc)
metadata = {} # Puedes agregar metadatos si quieres
return markdown, metadata
gr.Interface(
fn=convert,
inputs=[gr.File(label="Upload PDF", type="filepath")],
outputs=[gr.Text(label="Markdown crudo"), gr.JSON(label="Metadata")],
).launch()