Update app.py

app.py

@@ -6,22 +6,27 @@ built-in "Free OCR" prompt mode (plain text extraction, no markdown/
 grounding boxes). A "Markdown" mode is also included since it's a
 one-line prompt swap.
 
-This version runs on CPU only (no CUDA, no flash-attention). The model
-stays in bfloat16, matching hardcoded dtype casts inside its own custom
-code — PyTorch supports bfloat16 tensors on CPU, just without hardware
-acceleration, so it works but is slow. DeepSeek-OCR is a vision-language
-model, so expect anywhere from ~30s to a few minutes per image depending
-on resolution preset and CPU. Use the "Tiny" or "Small" resolution
-presets for faster (if lower-fidelity) results.
+This version runs on CPU only (no CUDA, no flash-attention). The model runs
+in float32 rather than bfloat16 — most CPUs have no fast bf16 GEMM kernel,
+so float32 is typically 2x+ faster here — and PyTorch's thread count is
+pinned to the number of available cores. DeepSeek-OCR is a vision-language
+model, so even with these speedups expect anywhere from ~15s to a couple
+of minutes per image depending on resolution preset and CPU. Use the
+"Tiny" or "Small" resolution presets, and the auto-crop option, for the
+fastest results.
 """
 
+import ast
 import os
 import re
 import shutil
 import tempfile
 
+import cv2
 import gradio as gr
+import numpy as np
 import torch
+from PIL import Image, ImageDraw, ImageFont
 from transformers import AutoModel, AutoTokenizer
 
 MODEL_NAME = "deepseek-ai/DeepSeek-OCR"
@@ -30,6 +35,21 @@ MODEL_NAME = "deepseek-ai/DeepSeek-OCR"
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 DEVICE = "cpu"
 
+# ---------------------------------------------------------------------------
+# CPU threading: containers sometimes leave PyTorch's thread count lower
+# than the number of cores actually available, which silently caps
+# performance. Pin it explicitly.
+# ---------------------------------------------------------------------------
+_num_cores = os.cpu_count() or 1
+torch.set_num_threads(_num_cores)
+try:
+    # Must be called before any inter-op parallel work has happened; if
+    # something already triggered that, this raises — safe to ignore.
+    torch.set_num_interop_threads(max(1, _num_cores // 2))
+except RuntimeError:
+    pass
+print(f"PyTorch using {_num_cores} threads for CPU inference.")
+
 # ---------------------------------------------------------------------------
 # Model loading (once, at startup)
 # ---------------------------------------------------------------------------
@@ -43,13 +63,6 @@ model = AutoModel.from_pretrained(
     trust_remote_code=True,
     use_safetensors=True,
 )
-# NOTE: don't cast to float32 here. DeepSeek-OCR's own forward pass hardcodes
-# casting image inputs to bfloat16 internally (independent of the model's
-# actual dtype), so the model's weights must also be bfloat16 or you get a
-# "Input type (BFloat16) and bias type (float) should be the same" crash.
-# PyTorch CPU does support bfloat16 tensors/ops (just without hardware
-# acceleration), so this works, it's just slow.
-model = model.eval().to(DEVICE).to(torch.bfloat16)
 
 # ---------------------------------------------------------------------------
 # DeepSeek-OCR's own `model.infer()` method has `.cuda()` calls hardcoded
@@ -59,10 +72,60 @@ model = model.eval().to(DEVICE).to(torch.bfloat16)
 # on CPU here, we monkey-patch `.cuda()` into a no-op that just returns the
 # tensor/module unchanged — this makes those hardcoded calls harmless instead
 # of failing.
+#
+# The same custom code also hardcodes casting some tensors to bfloat16 via
+# `.bfloat16()`, independent of the model's own weight dtype. Most CPUs
+# (anything short of newer server Xeons with AVX512_BF16/AMX, or recent
+# high-end AMD/Apple silicon) have no fast vectorized bf16 GEMM kernel, so
+# PyTorch silently falls back to a much slower generic path for bf16 on
+# those CPUs — often 2-4x slower than plain float32, which does have fast
+# native kernels everywhere. We keep the whole model in float32 and
+# monkey-patch `.bfloat16()` to a no-op (returning the tensor/module
+# converted to float32 instead) so those hardcoded casts don't undo it and
+# don't crash on a dtype mismatch either.
+#
+# NOTE: if you know for a fact your Space's CPU has native bf16
+# acceleration (uncommon on free/shared tiers), this patch may not help,
+# or could even hurt — comment out this block and use `.to(torch.bfloat16)`
+# below instead to compare.
 if not torch.cuda.is_available():
     torch.Tensor.cuda = lambda self, *args, **kwargs: self
     torch.nn.Module.cuda = lambda self, *args, **kwargs: self
 
+    def _redirect_bf16_to_fp32(args, kwargs):
+        args = list(args)
+        for i, a in enumerate(args):
+            if a is torch.bfloat16:
+                args[i] = torch.float32
+        if kwargs.get("dtype") is torch.bfloat16:
+            kwargs["dtype"] = torch.float32
+        return tuple(args), kwargs
+
+    # `.bfloat16()` patch (in case that exact method is called anywhere)...
+    torch.Tensor.bfloat16 = lambda self, *args, **kwargs: self.float()
+    torch.nn.Module.bfloat16 = lambda self, *args, **kwargs: self.float()
+
+    # ...but the actual crash showed the model's code uses
+    # `.to(torch.bfloat16)` / `.to(dtype=torch.bfloat16)` instead, which
+    # bypasses the patch above entirely (`.to()` is a different method).
+    # Intercept `.to()` itself and swap any bfloat16 request for float32,
+    # for both tensors and modules.
+    _orig_tensor_to = torch.Tensor.to
+    _orig_module_to = torch.nn.Module.to
+
+    def _patched_tensor_to(self, *args, **kwargs):
+        args, kwargs = _redirect_bf16_to_fp32(args, kwargs)
+        return _orig_tensor_to(self, *args, **kwargs)
+
+    def _patched_module_to(self, *args, **kwargs):
+        args, kwargs = _redirect_bf16_to_fp32(args, kwargs)
+        return _orig_module_to(self, *args, **kwargs)
+
+    torch.Tensor.to = _patched_tensor_to
+    torch.nn.Module.to = _patched_module_to
+
+model = model.eval().to(DEVICE).to(torch.float32)
+
 # Preset resolution modes exposed by the model's own infer() signature
 SIZE_PRESETS = {
     "Tiny (512, fastest)": dict(base_size=512, image_size=512, crop_mode=False),
@@ -75,21 +138,137 @@ SIZE_PRESETS = {
 TASK_PROMPTS = {
     "Free OCR (plain text)": "<image>\nFree OCR. ",
     "Markdown (structure + tables)": "<image>\n<|grounding|>Convert the document to markdown. ",
+    "Detect text regions (numbered boxes, no text output)": "<image>\n<|grounding|>OCR this image. ",
 }
 
+# The model emits boxes as <|ref|>label<|/ref|><|det|>[[x1,y1,x2,y2], ...]<|/det|>
+# with coordinates normalized to a 0-999 scale (independent of actual image
+# size), in the order the model reads them — so numbering them in the order
+# they appear already gives a natural left-to-right/top-to-bottom-ish
+# reading-order sequence.
+_REF_DET_RE = re.compile(
+    r"<\|ref\|>(.*?)<\|/ref\|>\s*<\|det\|>\s*(\[\[.*?\]\])\s*<\|/det\|>",
+    re.DOTALL,
+)
+
+
+def _load_font(size: int):
+    try:
+        return ImageFont.truetype(
+            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", size
+        )
+    except Exception:
+        return ImageFont.load_default()
+
+
+def parse_grounding_boxes(raw_text: str):
+    """
+    Extract [(label, [x1, y1, x2, y2]), ...] box groups from the model's raw
+    grounding output, in the order they appear. Coordinates are 0-999
+    normalized; the caller scales them to the actual image size. A single
+    <|ref|>/<|det|> pair can itself contain multiple boxes (e.g. a repeated
+    word) — those are flattened into separate entries in order too.
+    """
+    boxes = []
+    for label, coords_str in _REF_DET_RE.findall(raw_text):
+        try:
+            coord_list = ast.literal_eval(coords_str)
+        except (ValueError, SyntaxError):
+            continue
+        for coords in coord_list:
+            if len(coords) == 4:
+                boxes.append((label.strip(), coords))
+    return boxes
+
+
+def draw_numbered_boxes(image: Image.Image, boxes) -> Image.Image:
+    """Draw each box from parse_grounding_boxes() with a sequence number
+    (1, 2, 3, ...) in the order given — no recognized text is shown."""
+    img_out = image.convert("RGB").copy()
+    draw = ImageDraw.Draw(img_out)
+    w, h = img_out.size
+    font = _load_font(max(14, min(w, h) // 40))
+
+    for i, (_, (x1, y1, x2, y2)) in enumerate(boxes, start=1):
+        # Coordinates are normalized 0-999 regardless of actual image size.
+        px1, py1 = x1 / 999 * w, y1 / 999 * h
+        px2, py2 = x2 / 999 * w, y2 / 999 * h
 
-def run_ocr(image, task_choice, size_choice, progress=gr.Progress(track_tqdm=True)):
+        draw.rectangle([px1, py1, px2, py2], outline=(255, 0, 0), width=3)
+
+        label = str(i)
+        text_bbox = draw.textbbox((0, 0), label, font=font)
+        tw, th = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
+        tag_x0, tag_y0 = px1, max(0, py1 - th - 6)
+        draw.rectangle([tag_x0, tag_y0, tag_x0 + tw + 6, tag_y0 + th + 6],
+                       fill=(255, 0, 0))
+        draw.text((tag_x0 + 3, tag_y0 + 2), label, fill=(255, 255, 255), font=font)
+
+    return img_out
+
+
+def auto_crop_to_content(pil_img: Image.Image, pad_frac: float = 0.03) -> Image.Image:
+    """
+    Crop out large flat background regions (wall, table, frame margin, etc.)
+    so the model only has to encode the actual document/text area.
+
+    This is a cheap OpenCV threshold + bounding-box operation, not another
+    neural model — it costs a few milliseconds, versus the multi-second/
+    multi-tile cost DeepSeek-OCR would otherwise spend encoding background
+    pixels (especially in "Gundam" crop mode, where a large image gets
+    split into several 640x640 tiles). If nothing usable is found, the
+    original image is returned unchanged.
+    """
+    img = np.array(pil_img.convert("RGB"))
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+
+    # Otsu's method picks a threshold automatically; THRESH_BINARY_INV so
+    # content (usually darker than a plain background) shows up as
+    # foreground. Works well for "document/poster on a plain background"
+    # photos; less reliable for busy/textured backgrounds.
+    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+
+    # Close small gaps so text lines merge into solid blobs rather than
+    # hundreds of tiny disconnected contours.
+    kernel = np.ones((15, 15), np.uint8)
+    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
+
+    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return pil_img
+
+    all_points = np.vstack(contours)
+    x, y, w, h = cv2.boundingRect(all_points)
+
+    # Bail out if the detected region is basically the whole image already
+    # (nothing to gain) or suspiciously tiny (likely noise, not content).
+    img_area = img.shape[0] * img.shape[1]
+    box_area = w * h
+    if box_area < 0.02 * img_area or box_area > 0.98 * img_area:
+        return pil_img
+
+    pad_x, pad_y = int(w * pad_frac), int(h * pad_frac)
+    x0, y0 = max(0, x - pad_x), max(0, y - pad_y)
+    x1, y1 = min(img.shape[1], x + w + pad_x), min(img.shape[0], y + h + pad_y)
+    return pil_img.crop((x0, y0, x1, y1))
+
+
+def run_ocr(image, task_choice, size_choice, auto_crop, progress=gr.Progress(track_tqdm=True)):
     if image is None:
-        return "Please upload an image first."
+        return "Please upload an image first.", None
 
     prompt = TASK_PROMPTS[task_choice]
     preset = SIZE_PRESETS[size_choice]
 
+    image = image.convert("RGB")
+    if auto_crop:
+        image = auto_crop_to_content(image)
+
     # model.infer() reads the image from a file path and writes its result to
     # output_path. Point both at a temp dir.
     tmp_dir = tempfile.mkdtemp(prefix="deepseek_ocr_")
     image_path = os.path.join(tmp_dir, "input.png")
-    image.convert("RGB").save(image_path)
+    image.save(image_path)
 
     try:
         with torch.no_grad():
@@ -112,20 +291,26 @@ def run_ocr(image, task_choice, size_choice, progress=gr.Progress(track_tqdm=Tru
 
         result_file = os.path.join(tmp_dir, "result.mmd")
         if not os.path.exists(result_file):
-            return ("No output file was produced. Try a different resolution "
-                    "preset, or check the Space logs for errors.")
+            msg = ("No output file was produced. Try a different resolution "
+                   "preset, or check the Space logs for errors.")
+            return msg, None
 
         with open(result_file, "r", encoding="utf-8") as f:
-            text = f.read()
+            raw_text = f.read()
 
-        # Free OCR output is plain text already; Markdown/grounding output
-        # can contain <|ref|>...<|/ref|> and <|det|>[[...]]<|/det|> tags —
-        # strip those so the box shows clean readable text either way.
-        text = re.sub(r"<\|ref\|>.*?<\|/ref\|>", "", text, flags=re.DOTALL)
+        if task_choice == "Detect text regions (numbered boxes, no text output)":
+            boxes = parse_grounding_boxes(raw_text)
+            if not boxes:
+                return ("No text regions were detected in this image.", None)
+            annotated = draw_numbered_boxes(image, boxes)
+            return (f"Detected {len(boxes)} text region(s) — see numbered "
+                    f"boxes on the image.", annotated)
+
+        # Plain-text tasks: strip grounding markup if present and show text.
+        text = re.sub(r"<\|ref\|>.*?<\|/ref\|>", "", raw_text, flags=re.DOTALL)
         text = re.sub(r"<\|det\|>\[\[.*?\]\]<\|/det\|>", "", text, flags=re.DOTALL)
         text = text.strip()
-
-        return text if text else "No text was detected in this image."
+        return (text if text else "No text was detected in this image.", None)
     finally:
         shutil.rmtree(tmp_dir, ignore_errors=True)
 
@@ -155,17 +340,24 @@ with gr.Blocks(title="DeepSeek-OCR (CPU)") as demo:
                 value="Gundam (1024/640 crop, best for documents)",
                 label="Resolution mode (lower = faster on CPU)",
             )
+            auto_crop_checkbox = gr.Checkbox(
+                value=True,
+                label="Auto-crop background before OCR (faster on photos with a lot of empty space)",
+            )
             run_button = gr.Button("Run OCR", variant="primary")
 
         with gr.Column():
             output_text = gr.Textbox(
-                label="Extracted text", lines=25
+                label="Extracted text / status", lines=15
+            )
+            output_image = gr.Image(
+                label="Detected regions (numbered boxes)", visible=True
             )
 
     run_button.click(
         fn=run_ocr,
-        inputs=[image_input, task_dropdown, size_dropdown],
-        outputs=output_text,
+        inputs=[image_input, task_dropdown, size_dropdown, auto_crop_checkbox],
+        outputs=[output_text, output_image],
     )
 
 if __name__ == "__main__":