Update app.py

app.py

@@ -1,47 +1,52 @@
 import os
 import re
-import tempfile
 from io import BytesIO
-from typing import List, Tuple, Optional
+from typing import List, Tuple
 
 import gradio as gr
 import torch
 import numpy as np
-from PIL import Image, ImageDraw, ImageFont, ImageOps
+from PIL import Image, ImageDraw, ImageOps
 import fitz  # PyMuPDF
 
 from transformers import (
-    AutoProcessor,
+    TrOCRProcessor,
     VisionEncoderDecoderModel,
     BlipProcessor,
     BlipForConditionalGeneration,
 )
+from transformers.utils import logging as hf_logging
 
 # -------------------------
-# CPU-only setup
+# CPU-only, quieter logs
 # -------------------------
+hf_logging.set_verbosity_error()
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
 DEVICE = torch.device("cpu")
 torch.set_num_threads(int(os.getenv("TORCH_NUM_THREADS", "4")))
 
 TROCR_NAME = os.getenv("TROCR_MODEL", "microsoft/trocr-base-printed")
 BLIP_NAME = os.getenv("BLIP_MODEL", "Salesforce/blip-image-captioning-base")
 
+DEFAULT_DPI = 200
+MAX_SIDE = int(os.getenv("MAX_SIDE", "1600"))  # soft cap for CPU speed
+
 # -------------------------
 # Models (CPU)
 # -------------------------
-trocr_processor = AutoProcessor.from_pretrained(TROCR_NAME)
+trocr_processor = TrOCRProcessor.from_pretrained(TROCR_NAME)
 trocr_model = VisionEncoderDecoderModel.from_pretrained(TROCR_NAME).eval().to(DEVICE)
 
 blip_processor = BlipProcessor.from_pretrained(BLIP_NAME)
 blip_model = BlipForConditionalGeneration.from_pretrained(BLIP_NAME).eval().to(DEVICE)
 
 # -------------------------
-# Optional: pytesseract (for boxes on images)
+# Optional: Tesseract for image boxes
 # -------------------------
 def _try_import_tesseract():
     try:
         import pytesseract  # type: ignore
-        # Quick sanity check: version call triggers binary lookup
         _ = pytesseract.get_tesseract_version()
         return pytesseract
     except Exception:
@@ -49,44 +54,28 @@ def _try_import_tesseract():
 
 PYTESS = _try_import_tesseract()
 
-# -------------------------
-# UI / tasks
-# -------------------------
-TASKS = [
-    "OCR",
-    "Markdown",
-    "Locate",
-    "Describe",
-]
-
-DEFAULT_DPI = 200  # PDF render DPI
+TASKS = ["OCR", "Markdown", "Locate", "Describe"]
 
 
 # -------------------------
 # Helpers
 # -------------------------
-def _safe_font(size: int = 28):
-    candidates = [
-        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
-        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
-    ]
-    for p in candidates:
-        try:
-            if os.path.exists(p):
-                return ImageFont.truetype(p, size)
-        except Exception:
-            pass
-    return ImageFont.load_default()
-
-
 def _to_rgb(img: Image.Image) -> Image.Image:
     if img.mode in ("RGBA", "LA", "P"):
         img = img.convert("RGB")
-    return ImageOps.exif_transpose(img)
+    img = ImageOps.exif_transpose(img)
+
+    # Keep CPU inference reasonable
+    w, h = img.size
+    m = max(w, h)
+    if m > MAX_SIDE:
+        scale = MAX_SIDE / float(m)
+        img = img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
+    return img
 
 
 def _tokenize(s: str) -> List[str]:
-    return re.findall(r"[A-Za-zА-Яа-я0-9]+", s.lower())
+    return re.findall(r"[A-Za-zА-Яа-я0-9]+", (s or "").lower())
 
 
 def trocr_ocr(img: Image.Image) -> str:
@@ -96,7 +85,7 @@ def trocr_ocr(img: Image.Image) -> str:
     with torch.no_grad():
         ids = trocr_model.generate(pixel_values, max_new_tokens=256)
     text = trocr_processor.batch_decode(ids, skip_special_tokens=True)[0]
-    return text.strip()
+    return (text or "").strip()
 
 
 def blip_describe(img: Image.Image) -> str:
@@ -107,33 +96,25 @@ def blip_describe(img: Image.Image) -> str:
     return blip_processor.decode(out[0], skip_special_tokens=True).strip()
 
 
-def render_pdf_page(path: str, page_num: int, dpi: int = DEFAULT_DPI) -> Tuple[fitz.Document, fitz.Page, Image.Image, float]:
+def render_pdf_page(path: str, page_num: int, dpi: int = DEFAULT_DPI):
     doc = fitz.open(path)
-    page_idx = max(0, min(page_num - 1, len(doc) - 1))
+    page_idx = max(0, min(int(page_num) - 1, len(doc) - 1))
     page = doc.load_page(page_idx)
     zoom = dpi / 72.0
     pix = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
     img = Image.open(BytesIO(pix.tobytes("png")))
-    return doc, page, img, zoom
+    return doc, page, _to_rgb(img), zoom
 
 
 def pdf_has_text(page: fitz.Page) -> bool:
-    # words is empty for scanned pages
-    words = page.get_text("words")
-    return bool(words)
+    return bool(page.get_text("words"))
 
 
 def pdf_extract_text(page: fitz.Page) -> str:
-    txt = page.get_text("text") or ""
-    return txt.strip()
+    return (page.get_text("text") or "").strip()
 
 
 def pdf_to_markdown_simple(page: fitz.Page) -> str:
-    """
-    Lightweight markdown for selectable-text PDFs.
-    - Uses span sizes to guess headers.
-    - No heavy layout logic (keeps it stable and fast on CPU).
-    """
     data = page.get_text("dict")
     spans = []
     for b in data.get("blocks", []):
@@ -149,7 +130,7 @@ def pdf_to_markdown_simple(page: fitz.Page) -> str:
     h1_thr = med * 1.60
     h2_thr = med * 1.35
 
-    lines_out: List[str] = []
+    out_lines: List[str] = []
     for b in data.get("blocks", []):
         if b.get("type") != 0:
             continue
@@ -157,26 +138,22 @@ def pdf_to_markdown_simple(page: fitz.Page) -> str:
             parts = []
             sizes = []
             for sp in ln.get("spans", []):
-                t = (sp.get("text") or "")
-                if t.strip():
-                    parts.append(t.strip())
+                t = (sp.get("text") or "").strip()
+                if t:
+                    parts.append(t)
                     sizes.append(float(sp.get("size", 0.0)))
             if not parts:
                 continue
             line = " ".join(parts).strip()
             sz = max(sizes) if sizes else med
-
             if sz >= h1_thr:
-                lines_out.append("# " + line)
+                out_lines.append("# " + line)
             elif sz >= h2_thr:
-                lines_out.append("## " + line)
+                out_lines.append("## " + line)
             else:
-                lines_out.append(line)
-
-        lines_out.append("")  # paragraph break
-
-    md = "\n".join(lines_out).strip()
-    return md
+                out_lines.append(line)
+        out_lines.append("")
+    return "\n".join(out_lines).strip()
 
 
 def draw_rects(img: Image.Image, rects_px: List[Tuple[int, int, int, int]]) -> Image.Image:
@@ -192,23 +169,20 @@ def draw_rects(img: Image.Image, rects_px: List[Tuple[int, int, int, int]]) -> I
 
 
 def locate_in_pdf_words(page: fitz.Page, query: str) -> List[Tuple[float, float, float, float]]:
-    """
-    Returns list of rectangles in PDF coordinate space (points).
-    Uses exact word sequence match (token-based).
-    """
     q = _tokenize(query)
     if not q:
         return []
-
-    words = page.get_text("words")  # x0,y0,x1,y1,"word",block,line,wordno
+    words = page.get_text("words")
     if not words:
         return []
 
-    w_tokens = [_tokenize(w[4])[0] if _tokenize(w[4]) else "" for w in words]
-    rects: List[Tuple[float, float, float, float]] = []
+    w_tokens = []
+    for w in words:
+        toks = _tokenize(w[4])
+        w_tokens.append(toks[0] if toks else "")
 
-    n = len(w_tokens)
-    m = len(q)
+    rects = []
+    n, m = len(w_tokens), len(q)
     for i in range(0, n - m + 1):
         if w_tokens[i:i + m] == q:
             xs0 = [float(words[j][0]) for j in range(i, i + m)]
@@ -216,24 +190,17 @@ def locate_in_pdf_words(page: fitz.Page, query: str) -> List[Tuple[float, float,
             xs1 = [float(words[j][2]) for j in range(i, i + m)]
             ys1 = [float(words[j][3]) for j in range(i, i + m)]
             rects.append((min(xs0), min(ys0), max(xs1), max(ys1)))
-
     return rects
 
 
-def locate_in_image_tesseract(img: Image.Image, query: str) -> Tuple[List[Tuple[int, int, int, int]], str]:
-    """
-    Returns pixel-space rectangles for located phrase, plus a short status message.
-    If pytesseract is not available, returns empty list and message.
-    """
+def locate_in_image_tesseract(img: Image.Image, query: str):
     if PYTESS is None:
-        return [], "Tesseract not available: no boxes for images."
-
+        return [], "Tesseract not available."
     q = _tokenize(query)
     if not q:
         return [], "Empty query."
 
     img = _to_rgb(img)
-    # Use data dict so it works consistently
     data = PYTESS.image_to_data(img, output_type=PYTESS.Output.DICT)
 
     texts = data.get("text", [])
@@ -249,148 +216,131 @@ def locate_in_image_tesseract(img: Image.Image, query: str) -> Tuple[List[Tuple[
         t = (t or "").strip()
         if not t:
             continue
-        tok = _tokenize(t)
-        if not tok:
+        toks = _tokenize(t)
+        if not toks:
             continue
-        # Keep only "reasonable" confidence if numeric
         try:
             c = float(conf[i])
             if c < 0:
                 continue
         except Exception:
             pass
-
-        tokens.append(tok[0])
+        tokens.append(toks[0])
         boxes.append((int(left[i]), int(top[i]), int(left[i] + width[i]), int(top[i] + height[i])))
 
-    rects: List[Tuple[int, int, int, int]] = []
-    n = len(tokens)
-    m = len(q)
+    rects_px = []
+    n, m = len(tokens), len(q)
     for i in range(0, n - m + 1):
         if tokens[i:i + m] == q:
             xs0 = [boxes[j][0] for j in range(i, i + m)]
             ys0 = [boxes[j][1] for j in range(i, i + m)]
             xs1 = [boxes[j][2] for j in range(i, i + m)]
             ys1 = [boxes[j][3] for j in range(i, i + m)]
-            rects.append((min(xs0), min(ys0), max(xs1), max(ys1)))
+            rects_px.append((min(xs0), min(ys0), max(xs1), max(ys1)))
 
-    if not rects:
-        return [], "Not found."
-    return rects, "Found."
+    return rects_px, ("Found." if rects_px else "Not found.")
 
 
-def as_markdown_block(text: str) -> str:
-    if not text.strip():
-        return ""
-    return "```text\n" + text.strip() + "\n```"
+def as_text_block(s: str) -> str:
+    s = (s or "").strip()
+    return s if s else ""
 
 
 # -------------------------
-# Main run
+# Core processing
 # -------------------------
-def process(path: str, task: str, page_num: int, query: str):
-    if not path:
-        return "Upload a file.", "", None
+def process(file_path: str, task: str, page_num: int, query: str):
+    if not file_path:
+        return "Upload a file.", "", None, None
 
-    ext = os.path.splitext(path)[1].lower()
+    ext = os.path.splitext(file_path)[1].lower()
 
-    # ---------- PDF ----------
+    # PDF
     if ext == ".pdf":
-        doc, page, page_img, zoom = render_pdf_page(path, page_num, dpi=DEFAULT_DPI)
+        doc, page, page_img, zoom = render_pdf_page(file_path, page_num, dpi=DEFAULT_DPI)
         try:
+            preview = page_img
+
             if task == "Describe":
-                caption = blip_describe(page_img)
-                return caption, as_markdown_block(caption), None
+                cap = blip_describe(page_img)
+                return cap, cap, None, preview
 
             if task == "OCR":
-                if pdf_has_text(page):
-                    txt = pdf_extract_text(page)
-                else:
-                    txt = trocr_ocr(page_img)
-                return txt, as_markdown_block(txt), None
+                txt = pdf_extract_text(page) if pdf_has_text(page) else trocr_ocr(page_img)
+                return txt, txt, None, preview
 
             if task == "Markdown":
                 if pdf_has_text(page):
                     md = pdf_to_markdown_simple(page)
                     if not md:
-                        txt = pdf_extract_text(page)
-                        md = as_markdown_block(txt)
+                        md = pdf_extract_text(page)
                 else:
-                    txt = trocr_ocr(page_img)
-                    md = as_markdown_block(txt)
-                return md, md, None
+                    md = trocr_ocr(page_img)
+                return md, md, None, preview
 
             if task == "Locate":
-                if not query.strip():
-                    return "Enter text to locate.", "", page_img
+                if not (query or "").strip():
+                    return "Enter query.", "", preview, preview
 
-                # 1) Prefer precise PDF word boxes (selectable text)
+                # selectable-text PDF: precise boxes
                 rects_pdf = locate_in_pdf_words(page, query)
                 if rects_pdf:
-                    # Convert PDF points -> pixels using same render zoom
-                    rects_px = []
-                    for (x0, y0, x1, y1) in rects_pdf:
-                        rects_px.append((int(x0 * zoom), int(y0 * zoom), int(x1 * zoom), int(y1 * zoom)))
+                    rects_px = [(int(x0 * zoom), int(y0 * zoom), int(x1 * zoom), int(y1 * zoom)) for x0, y0, x1, y1 in rects_pdf]
                     boxed = draw_rects(page_img, rects_px)
-                    return "Found.", "", boxed
+                    return "Found.", "", boxed, preview
 
-                # 2) Fallback: if scanned page, try tesseract boxes on rendered image
+                # fallback: render + tesseract
                 rects_px, msg = locate_in_image_tesseract(page_img, query)
                 boxed = draw_rects(page_img, rects_px) if rects_px else page_img
-                return msg, "", boxed
+                return msg, "", boxed, preview
 
-            return "Unknown task.", "", None
+            return "Unknown task.", "", None, preview
         finally:
             doc.close()
 
-    # ---------- Image ----------
-    img = _to_rgb(Image.open(path))
+    # Image
+    img = _to_rgb(Image.open(file_path))
+    preview = img
 
     if task == "Describe":
-        caption = blip_describe(img)
-        return caption, as_markdown_block(caption), None
+        cap = blip_describe(img)
+        return cap, cap, None, preview
 
     if task == "OCR":
         txt = trocr_ocr(img)
-        return txt, as_markdown_block(txt), None
+        return txt, txt, None, preview
 
     if task == "Markdown":
-        txt = trocr_ocr(img)
-        md = as_markdown_block(txt)
-        return md, md, None
+        md = trocr_ocr(img)
+        return md, md, None, preview
 
     if task == "Locate":
-        if not query.strip():
-            return "Enter text to locate.", "", img
-
+        if not (query or "").strip():
+            return "Enter query.", "", img, preview
         rects_px, msg = locate_in_image_tesseract(img, query)
         boxed = draw_rects(img, rects_px) if rects_px else img
-        return msg, "", boxed
+        return msg, "", boxed, preview
 
-    return "Unknown task.", "", None
+    return "Unknown task.", "", None, preview
 
 
 # -------------------------
-# UI helpers
+# UI wiring
 # -------------------------
-def update_page_selector(file_path: str):
+def update_page_ui(file_path: str):
     if not file_path:
-        return gr.update(visible=False), gr.update(value=None)
+        return gr.update(visible=False), None
 
     ext = os.path.splitext(file_path)[1].lower()
     if ext != ".pdf":
-        return gr.update(visible=False), gr.update(value=_to_rgb(Image.open(file_path)))
+        return gr.update(visible=False), _to_rgb(Image.open(file_path))
 
     doc = fitz.open(file_path)
-    pages = len(doc)
+    pages = max(1, len(doc))
     doc.close()
 
-    # Show first page preview
     _, _, img, _ = render_pdf_page(file_path, 1, dpi=DEFAULT_DPI)
-    return (
-        gr.update(visible=True, minimum=1, maximum=max(1, pages), value=1),
-        gr.update(value=img),
-    )
+    return gr.update(visible=True, minimum=1, maximum=pages, value=1), img
 
 
 def update_preview(file_path: str, page_num: int):
@@ -408,43 +358,37 @@ def toggle_query(task: str):
 
 
 # -------------------------
-# Build app (minimal style)
+# Minimal UI style
 # -------------------------
-theme = gr.themes.Base(
-    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
+theme = gr.themes.Monochrome(
+    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"]
 )
 
 with gr.Blocks(theme=theme, title="Doc Tool (CPU)") as demo:
     with gr.Row():
         with gr.Column(scale=1, min_width=320):
             file_in = gr.File(label="File", file_types=["image", ".pdf"], type="filepath")
-            page_num = gr.Slider(label="Page", minimum=1, maximum=1, value=1, step=1, visible=False)
+            page = gr.Slider(label="Page", minimum=1, maximum=1, value=1, step=1, visible=False)
             task = gr.Dropdown(label="Task", choices=TASKS, value="OCR")
-            query = gr.Textbox(label="Query", visible=False, placeholder="Text to locate")
-
+            query = gr.Textbox(label="Query", placeholder="Text to locate", visible=False)
             run_btn = gr.Button("Run", variant="primary")
 
         with gr.Column(scale=2):
-            preview = gr.Image(label="Preview", type="pil", height=360)
-            out_text = gr.Textbox(label="Output", lines=10)
-            out_md = gr.Markdown()
-
-            out_boxes = gr.Image(label="Boxes", type="pil", height=360)
+            with gr.Row():
+                preview = gr.Image(label="Preview", type="pil", height=320)
+                boxes = gr.Image(label="Boxes", type="pil", height=320)
+            out = gr.Textbox(label="Output", lines=10)
 
-    file_in.change(update_page_selector, inputs=[file_in], outputs=[page_num, preview])
-    page_num.change(update_preview, inputs=[file_in, page_num], outputs=[preview])
+    file_in.change(update_page_ui, inputs=[file_in], outputs=[page, preview])
+    page.change(update_preview, inputs=[file_in, page], outputs=[preview])
     task.change(toggle_query, inputs=[task], outputs=[query])
 
-    def on_run(file_path, task_name, page, q):
-        text, md, boxed = process(file_path, task_name, int(page), q or "")
-        return text, md, boxed
+    def on_run(fp, t, p, q):
+        text, _, boxed, prev = process(fp, t, int(p), q or "")
+        # keep preview stable; boxes only when relevant
+        return prev, boxed, as_text_block(text)
 
-    run_btn.click(
-        on_run,
-        inputs=[file_in, task, page_num, query],
-        outputs=[out_text, out_md, out_boxes],
-    )
+    run_btn.click(on_run, inputs=[file_in, task, page, query], outputs=[preview, boxes, out])
 
 if __name__ == "__main__":
-    # Disable SSR to avoid extra startup noise
     demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)