Update app.py

app.py

@@ -1,156 +1,132 @@
-# raw_paddleocr.py
-# Standalone raw-text extractor using PaddleOCR (no changes to your app).
-# Modes:
-#   OCR_RAW_MODE = "block" (default) | "paragraph" | "lines"
-#   OCR_CONF_THRESHOLD = 0.0..1.0  (default 0.0)
-#   OCR_LANG = "en" (default) or other PaddleOCR langs like "ar", "en_number"
-#   OCR_USE_GPU = "true" | "false" (default "false")
-
 import os
+import io
 import sys
-from typing import List, Tuple, Dict, Any
+import json
+import traceback
+from typing import List, Tuple
 
 import numpy as np
 from PIL import Image
 import fitz  # PyMuPDF
 import cv2
-
+import gradio as gr
 from paddleocr import PaddleOCR
 
-# -------- Config (env-driven) ----------
-LANG = os.getenv("OCR_LANG", "en")
+# --------- Config knobs (safe defaults) ----------
+LANG = os.getenv("OCR_LANG", "en")          # e.g., "en", "ar", "en_number", "en_PP-OCRv3"
 USE_GPU = os.getenv("OCR_USE_GPU", "false").lower() == "true"
-CLS = True
-CONF_THRESHOLD = float(os.getenv("OCR_CONF_THRESHOLD", "0.0"))
-
-RAW_MODE = os.getenv("OCR_RAW_MODE", "block")  # "block" | "paragraph" | "lines"
-LINE_GAP_RATIO = float(os.getenv("OCR_LINE_GAP_RATIO", "0.6"))
+DET = os.getenv("OCR_DET_MODEL", "ch_PP-OCRv4_det")
+REC = os.getenv("OCR_REC_MODEL", "en_PP-OCRv4")
+CLS = True                                  # angle classification
+CONF_THRESHOLD = float(os.getenv("OCR_CONF_THRESHOLD", "0.0"))  # 0.0 → keep everything
 
-# -------- Init OCR once ---------------
+# Initialize once (download models once, reuse across requests)
+# Tip: If you want Arabic/English mixed, set LANG="ar" or "en" variants per PaddleOCR docs
 OCR = PaddleOCR(
     use_angle_cls=CLS,
     lang=LANG,
     use_gpu=USE_GPU,
-    det_model_dir=None,
-    rec_model_dir=None,
+    det_model_dir=None,   # use default
+    rec_model_dir=None,   # use default
     show_log=False
 )
 
-# -------- Utils -----------------------
 def _pil_to_cv(img: Image.Image) -> np.ndarray:
+    """PIL RGB -> OpenCV BGR ndarray"""
     return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
 
-def _bbox_center(bbox):
-    xs = [p[0] for p in bbox]; ys = [p[1] for p in bbox]
-    return (sum(xs) / 4.0, sum(ys) / 4.0)
-
-def read_image(path: str) -> Image.Image:
-    with Image.open(path) as im:
+def ocr_image(pil_img: Image.Image) -> List[Tuple[str, float]]:
+    """
+    Run OCR on a PIL image and return list of (text, confidence).
+    """
+    img_cv = _pil_to_cv(pil_img)
+    result = OCR.ocr(img_cv, cls=CLS)
+    lines: List[Tuple[str, float]] = []
+    if not result:
+        return lines
+    # PaddleOCR returns a list per image; each item has [ [box, (text, conf)], ... ]
+    for line in result[0]:
+        txt = line[1][0]
+        conf = float(line[1][1])
+        if conf >= CONF_THRESHOLD:
+            lines.append((txt, conf))
+    return lines
+
+def read_image(filepath: str) -> Image.Image:
+    """
+    Open an image robustly via PIL (also handles TIFF, JPG, PNG).
+    """
+    with Image.open(filepath) as im:
         return im.convert("RGB")
 
-def read_pdf_pages(path: str) -> List[Image.Image]:
+def read_pdf_pages(filepath: str) -> List[Image.Image]:
+    """
+    Render each PDF page to a PIL image (RGB) using PyMuPDF.
+    """
     pages: List[Image.Image] = []
-    with fitz.open(path) as doc:
+    with fitz.open(filepath) as doc:
         for page in doc:
-            mat = fitz.Matrix(2, 2)  # scale up for better OCR
+            # Render with a scale factor for better OCR accuracy
+            mat = fitz.Matrix(2, 2)  # 2x upscaling
             pix = page.get_pixmap(matrix=mat, alpha=False)
             img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
             pages.append(img)
     return pages
 
-# -------- Core OCR --------------------
-def ocr_tokens(pil_img: Image.Image) -> List[Dict[str, Any]]:
-    img_cv = _pil_to_cv(pil_img)
-    result = OCR.ocr(img_cv, cls=CLS)
-    tokens = []
-    if not result:
-        return tokens
-    for box, (txt, conf) in result[0]:
-        conf = float(conf)
-        if not txt or conf < CONF_THRESHOLD:
-            continue
-        cx, cy = _bbox_center(box)
-        ys = [p[1] for p in box]
-        h = max(ys) - min(ys) + 1e-6
-        tokens.append({"text": txt.strip(), "conf": conf, "bbox": box, "cx": cx, "cy": cy, "h": h})
-    return tokens
-
-def _sort_reading_order(tokens: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-    return sorted(tokens, key=lambda t: (round(t["cy"], 1), t["cx"]))
-
-def _group_paragraphs(sorted_tokens: List[Dict[str, Any]], gap_ratio: float = 0.6):
-    if not sorted_tokens:
-        return []
-    heights = sorted(t["h"] for t in sorted_tokens)
-    median_h = heights[len(heights)//2] or 1.0
-
-    paras, cur = [], [sorted_tokens[0]]
-    for prev, cur_tok in zip(sorted_tokens, sorted_tokens[1:]):
-        vertical_gap = cur_tok["cy"] - prev["cy"]
-        if vertical_gap > gap_ratio * median_h:
-            paras.append(cur)
-            cur = [cur_tok]
-        else:
-            cur.append(cur_tok)
-    paras.append(cur)
-    return paras
-
-def _post_clean(text: str) -> str:
-    text = " ".join(text.split())
-    text = text.replace("- ", "")
-    return text
-
-def tokens_to_text(tokens: List[Dict[str, Any]], mode: str = "block", gap_ratio: float = 0.6) -> str:
-    if not tokens:
-        return ""
-    tokens = _sort_reading_order(tokens)
-
-    if mode == "block":
-        return _post_clean(" ".join(t["text"] for t in tokens))
-
-    if mode == "paragraph":
-        paras = _group_paragraphs(tokens, gap_ratio=gap_ratio)
-        chunks = [_post_clean(" ".join(t["text"] for t in p)) for p in paras]
-        return "\n\n".join(c for c in chunks if c)
-
-    # lines
-    lines, current = [], [tokens[0]]
-    for prev, cur_tok in zip(tokens, tokens[1:]):
-        same_line = abs(cur_tok["cy"] - prev["cy"]) <= 0.35 * max(prev["h"], cur_tok["h"])
-        if same_line:
-            current.append(cur_tok)
-        else:
-            lines.append(current)
-            current = [cur_tok]
-    lines.append(current)
-    line_texts = [_post_clean(" ".join(t["text"] for t in row)) for row in lines]
-    return "\n".join(l for l in line_texts if l)
-
-def extract_raw_text(path: str) -> str:
-    lower = path.lower()
+def extract_text_from_file(filepath: str) -> str:
+    """
+    Dispatch by file type; return plain text.
+    """
+    lower = filepath.lower()
     if lower.endswith(".pdf"):
-        pages = read_pdf_pages(path)
+        pages = read_pdf_pages(filepath)
+        all_text: List[str] = []
+        for i, pil_img in enumerate(pages, start=1):
+            lines = ocr_image(pil_img)
+            page_text = "\n".join([t for t, _ in lines])
+            # Add a page header for clarity on multi-page docs
+            all_text.append(f"--- Page {i} ---\n{page_text}".strip())
+        return "\n\n".join([s for s in all_text if s])
     elif lower.endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp")):
-        pages = [read_image(path)]
+        img = read_image(filepath)
+        lines = ocr_image(img)
+        return "\n".join([t for t, _ in lines]).strip()
     else:
-        raise ValueError("Unsupported file type. Provide an image or PDF.")
-
-    outputs = []
-    for pil_img in pages:
-        toks = ocr_tokens(pil_img)
-        outputs.append(tokens_to_text(toks, mode=RAW_MODE, gap_ratio=LINE_GAP_RATIO))
+        raise ValueError("Unsupported file type. Please upload an image (PNG/JPG/TIFF/WEBP/BMP) or a PDF.")
+
+def infer(file_obj) -> str:
+    try:
+        if file_obj is None:
+            return "No file uploaded."
+        filepath = file_obj.name if hasattr(file_obj, "name") else str(file_obj)
+        text = extract_text_from_file(filepath)
+        # 🔊 Console telemetry: dump raw text to terminal
+        print("\n================ OCR RAW TEXT ================\n")
+        print(text)
+        print("\n==================== END =====================\n", flush=True)
+        return text or "[No text detected]"
+    except Exception as e:
+        traceback.print_exc()
+        return f"Error during OCR: {e}"
+
+# ------------- Gradio UI ----------------
+TITLE = "PaddleOCR Text Extractor (Images & PDFs)"
+DESC = (
+    "Upload an image or PDF. The app runs PaddleOCR (PP-OCRv4 pipeline) and returns plain text. "
+    "Set `OCR_LANG`, `OCR_USE_GPU`, and `OCR_CONF_THRESHOLD` as env vars to tune."
+)
 
-    text = "\n\n".join(o for o in outputs if o).strip()
-    return text or "[No text detected]"
+with gr.Blocks(title=TITLE) as demo:
+    gr.Markdown(f"# {TITLE}\n{DESC}")
+    with gr.Row():
+        file_in = gr.File(label="Upload Image or PDF", file_count="single", file_types=["image", ".pdf"])
+    out = gr.Textbox(label="Extracted Text", lines=25, show_copy_button=True)
+    run_btn = gr.Button("Run OCR", variant="primary")
 
-# -------- CLI -------------------------
-def main():
-    if len(sys.argv) < 2:
-        print("Usage: python raw_paddleocr.py <path-to-image-or-pdf>")
-        sys.exit(2)
-    path = sys.argv[1]
-    out = extract_raw_text(path)
-    print(out)
+    run_btn.click(fn=infer, inputs=[file_in], outputs=[out])
+    # Also trigger on file change for convenience
+    file_in.change(fn=infer, inputs=[file_in], outputs=[out])
 
 if __name__ == "__main__":
-    main()
+    # Tip: Set server_name="0.0.0.0" for containers; share=True for quick external testing
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)