Update app.py

app.py

@@ -1,115 +1,197 @@
-import fitz  # PyMuPDF
-import pytesseract
-import io
 import os
+import unicodedata
+import fitz
 from PIL import Image
 import gradio as gr
-import tempfile
-import re
+import numpy as np
+import cv2
+from dotenv import load_dotenv
+import easyocr
+import pytesseract
+
+load_dotenv()
+
+reader = easyocr.Reader(['es', 'en'])
+
+def clean_text(text):
+    text = unicodedata.normalize("NFC", text)
+    lines = text.splitlines()
+    cleaned_lines = [line.strip() for line in lines if line.strip()]
+    return "\n".join(cleaned_lines)
+
+def clean_ocr_lines(text):
+    lines = text.splitlines()
+    cleaned = []
+    for line in lines:
+        line = line.strip()
+        if line:
+            line = " ".join(line.split())
+            cleaned.append(line)
+    return "\n".join(cleaned)
+
+def preprocess_for_ocr(pil_image):
+    gray = pil_image.convert('L')
+    np_img = np.array(gray)
+    try:
+        from skimage.filters import threshold_sauvola
+        window_size = 25
+        thresh_sauvola = threshold_sauvola(np_img, window_size=window_size)
+        binary = (np_img > thresh_sauvola).astype("uint8") * 255
+    except:
+        binary = cv2.adaptiveThreshold(np_img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+                                       cv2.THRESH_BINARY, 31, 10)
+    return Image.fromarray(binary)
 
-def extract_text_from_pdf(pdf_path):
-    text_output = []
+def run_easyocr(image_path):
+    img = Image.open(image_path)
+    img = preprocess_for_ocr(img)
+    img.save(image_path)
+    results = reader.readtext(image_path, detail=0, paragraph=False, decoder='greedy')
+    text = "\n".join(results)
+    return clean_ocr_lines(text)
+
+def run_tesseract_ocr(pil_image):
+    pil_image = preprocess_for_ocr(pil_image)
+    config = '--oem 3 --psm 6 -l spa+eng'
+    text = pytesseract.image_to_string(pil_image, config=config)
+    return clean_ocr_lines(text)
+
+def extract_embedded_images(page, page_number, seen_xrefs):
     image_paths = []
+    blocks = []
+    for img_index, img in enumerate(page.get_images(full=True)):
+        xref = img[0]
+        if xref in seen_xrefs:
+            continue
+        seen_xrefs.add(xref)
+        base_image = page.parent.extract_image(xref)
+        image_bytes = base_image["image"]
+        ext = base_image["ext"]
+        image_path = f"/tmp/embedded_p{page_number + 1}_{img_index + 1}.{ext}"
+        with open(image_path, "wb") as f:
+            f.write(image_bytes)
+        image_paths.append(image_path)
+        blocks.append(f"![Imagen_Embedded]({image_path})\n")
+    return blocks, image_paths
+
+def extract_visual_regions(image, page_number):
+    results = []
+    np_img = np.array(image.convert("RGB"))
+    gray = cv2.cvtColor(np_img, cv2.COLOR_RGB2GRAY)
+    _, binary = cv2.threshold(gray, 220, 255, cv2.THRESH_BINARY_INV)
+    closed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15)))
+    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(closed, connectivity=8)
+
+    for i in range(1, num_labels):
+        x, y, w, h, area = stats[i]
+        if area > 5000 and h > 50 and w > 50 and 0.3 < (w / float(h)) < 3.5:
+            bbox = (x, y, x + w, y + h)
+            crop = image.crop(bbox)
+            crop_path = f"/tmp/visual_crop_p{page_number + 1}_{i}.jpg"
+            crop.save(crop_path)
+            text_crop = run_tesseract_ocr(crop)
+            word_count = len(text_crop.split())
+            if 2 < word_count < 20:
+                results.append(crop_path)
+    return results
+
+def is_scanned_page(page):
+    text = page.get_text("text")
+    return not text or len(text.strip()) < 30
+
+def process_document(input_file):
+    if not input_file:
+        return None, "No file uploaded", None
 
-    with fitz.open(pdf_path) as doc:
-        for page_num, page in enumerate(doc):
-            # Render page to an image
+    temp_path = input_file.name
+    ext = os.path.splitext(temp_path)[-1].lower()
+    markdown_output = ""
+    all_images = []
+    seen_xrefs = set()
+
+    if ext in [".png", ".jpg", ".jpeg"]:
+        image = Image.open(temp_path)
+        text = run_tesseract_ocr(image)
+        markdown_output += f"## Resultado OCR\n\n{clean_text(text)}\n"
+        return markdown_output, [], None
+
+    doc = fitz.open(temp_path)
+    for i, page in enumerate(doc):
+        markdown_output += f"\n## Página {i + 1}\n\n"
+        text_dict = page.get_text("dict")
+        lines = []
+        for block in text_dict["blocks"]:
+            if "lines" in block:
+                for l in block["lines"]:
+                    line_parts = [span["text"].strip() for span in l["spans"] if span["text"].strip()]
+                    if line_parts:
+                        lines.append(" ".join(line_parts))
+                lines.append("")
+        text = "\n".join(lines).strip()
+
+        if not is_scanned_page(page):
+            markdown_output += f"{clean_text(text)}\n"
+        else:
             pix = page.get_pixmap(dpi=300)
-            image_path = f"/tmp/ocr_page_{page_num + 1}.jpg"
-            pix.save(image_path)
-            image_paths.append(image_path)
-
-            # Run OCR on the image
-            image = Image.open(image_path)
-            raw_text = pytesseract.image_to_string(image, lang='spa')
-
-            # Basic clean-up: remove gibberish if most lines are bad
-            def is_gibberish(text):
-                lines = text.splitlines()
-                bad_lines = [line for line in lines if len(re.findall(r'[a-zA-Z]', line)) < 5]
-                return len(bad_lines) / max(1, len(lines)) > 0.4
-
-            if is_gibberish(raw_text):
-                raw_text = ""
-
-            text_output.append({
-                "page_num": page_num + 1,
-                "image_path": image_path,
-                "text": raw_text
-            })
-
-    return text_output
-
-def generate_markdown(text_data, extra_image_path):
-    md = ""
-    for page in text_data:
-        md += f"## Página {page['page_num']}\n\n"
-        md += f"![Pagina Escaneada]({page['image_path']})\n\n"
-        md += ("### Detalle del error de carga\n\n"
-               "**Marca / Modelo:** VAG  \n"
-               "**Año:** 2014  \n"
-               "**Código de Motor:** EV/híbrido\n\n"
-               "**Síntoma / Código de Falla:**\n\n"
-               "> La carga de la batería de alto voltaje se interrumpe al cabo de aproximadamente 1 minuto. Sin embargo, se puede cargar cuando el automóvil está fuera de línea.  \n"
-               "> **No se almacenan códigos de falla.**\n\n"
-               "---\n\n"
-               "### Posible causa\n\n"
-               "El cliente ha instalado una aplicación de terceros que interrumpe la carga.  \n"
-               "Revisar el smartphone del cliente, que actúa como Master del coche.  \n"
-               "Las aplicaciones que estén conectadas al coche deben estar desconectadas.\n\n"
-               "Algunas aplicaciones pueden ajustar la carga para que consuma energía en los momentos en que la energía es más barata (por ejemplo, de noche o con energía solar).\n\n"
-               "Cuando la aplicación determina que no es beneficioso cargar, interrumpe el proceso.\n\n"
-               "---\n\n"
-               "### Solución sugerida\n\n"
-               "1. Poner el coche en **modo Offline** desde la función de taller.  \n"
-               "   El símbolo del globo cambiará de color:\n"
-               "   - **Globo gris** = modo sin conexión  \n"
-               "   - **Globo blanco** = modo online\n\n"
-               "2. Si al estar Offline el coche carga normalmente, es señal de que la aplicación es la causa.\n\n"
-               "3. **No basta con desinstalar la aplicación**:  \n"
-               "   Se debe **desvincular el coche** de ella por completo.\n\n"
-               "---\n\n"
-               "### Reinstalación (opcional)\n\n"
-               "El cliente puede optar por eliminar y reinstalar la aplicación para probar si una nueva conexión resuelve el problema.\n\n"
-               "---\n\n"
-               "### Aplicaciones conocidas que causan este problema\n\n"
-               "- Aplicación de coche eléctrico  \n"
-               "- Evcc  \n"
-               "- gridio  \n"
-               "- Github WeConnect-cli  \n"
-               "- tronidad  \n"
-               "- Elli Naturstrom  \n\n")
-
-        md += f"## Imagen relevante\n\n"
-        md += f"![Indicadores de carga e interfaz de enchufe]({extra_image_path})\n\n"
-        md += "---\n\n"
-    return md
-
-def ocr_app(file, extra_image):
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
-        tmp_file.write(file.read())
-        pdf_path = tmp_file.name
-
-    text_data = extract_text_from_pdf(pdf_path)
-    markdown_result = generate_markdown(text_data, extra_image.name)
-
-    output_md_path = "/tmp/resultado.md"
-    with open(output_md_path, "w") as f:
-        f.write(markdown_result)
-
-    return markdown_result, output_md_path
-
-demo = gr.Interface(
-    fn=ocr_app,
-    inputs=[
-        gr.File(label="Sube tu PDF", file_types=[".pdf"]),
-        gr.File(label="Imagen correcta (solo una)", file_types=[".png", ".jpg", ".jpeg"])
-    ],
-    outputs=[
-        gr.Markdown(label="Texto Extraído"),
-        gr.File(label="Descargar Markdown")
-    ],
-    title="OCR PDF - Extracción Limpia"
-)
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            image_path = f"/tmp/ocr_page_{i + 1}.jpg"
+            img.save(image_path)
+            all_images.append(image_path)
+            markdown_output += f"![Pagina_Scaneada]({image_path})\n\n"
+            ocr_text = run_tesseract_ocr(img)
+            markdown_output += f"{clean_text(ocr_text)}\n"
+            crops = extract_visual_regions(img, i)
+            for crop_path in crops:
+                all_images.append(crop_path)
+                markdown_output += f"![Region_Detectada]({crop_path})\n"
+
+        blocks, embedded_images = extract_embedded_images(page, i, seen_xrefs)
+        for block in blocks:
+            markdown_output += block
+        all_images.extend(embedded_images)
+        markdown_output += "\n---\n\n"
+
+    markdown_path = "/tmp/resultado.md"
+    with open(markdown_path, "w", encoding="utf-8") as f:
+        f.write(markdown_output)
+
+    return markdown_output.strip(), all_images, markdown_path
+
+# UI
+
+theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="rose", neutral_hue="stone")
+
+with gr.Blocks(theme=theme) as demo:
+    gr.Markdown("# OCR Preciso + Extracción Inteligente de Imágenes del PDF")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_file = gr.File(label="Sube PDF o Imagen", file_types=[".pdf", ".png", ".jpg", ".jpeg"])
+            run_button = gr.Button("Ejecutar OCR")
+        with gr.Column(scale=2):
+            markdown_output = gr.Textbox(
+                label="Markdown Generado",
+                lines=25,
+                max_lines=1000,
+                interactive=True,
+                elem_id="markdown_scrollbox"
+            )
+            gallery_output = gr.Gallery(label="Imágenes Extraídas", type="file")
+            download_md = gr.File(label="Descargar Markdown")
+
+    run_button.click(
+        fn=process_document,
+        inputs=[input_file],
+        outputs=[markdown_output, gallery_output, download_md]
+    )
+
+demo.css = """
+#markdown_scrollbox textarea {
+    overflow-y: auto !important;
+    max-height: 600px;
+    resize: vertical;
+    font-family: monospace;
+}
+"""
 
 demo.launch()