Update app.py

app.py

@@ -1,167 +1,221 @@
 """
-Eagle Eye — ZeroGPU OCR API Service (Hugging Face Space)
-Model: Qwen/Qwen2.5-VL-3B-Instruct  (fast, small, accurate)
+Eagle Eye — ZeroGPU OCR Service  v3 (YOLO + Qwen2.5-VL-7B)
+=============================================================
+Detection : comictextdetector.pt  (manga-image-translator YOLO model)
+OCR       : Qwen/Qwen2.5-VL-7B-Instruct
 
-Speed optimizations:
-  - True GPU batching: multiple images in one forward pass
-  - Greedy decoding (no sampling)
-  - max_new_tokens=128 (enough for speech bubbles)
-  - float16 weights
-  - Model loaded once and cached globally
-
-Upload this file + requirements.txt to your HF Space (ZeroGPU enabled).
+Input  : JSON list of base64-encoded full-page PNG images
+Output : JSON list of text strings (one per page, bubbles separated by \\n)
 """
 
-import gradio as gr
-import spaces
-import json
 import base64
 import io
+import json
 import traceback
+
+import gradio as gr
+import numpy as np
+import spaces
 from PIL import Image
 
-MODEL = None
-PROCESSOR = None
-DEVICE = None
+# ── Global model cache (loaded once, kept alive between ZeroGPU calls) ─────
 
+_YOLO = None        # ultralytics.YOLO  |  "fallback"
+_OCR_MODEL = None
+_OCR_PROC = None
+_DEVICE = None
 
-def load_model():
-    global MODEL, PROCESSOR, DEVICE
-    if MODEL is not None:
-        return MODEL, PROCESSOR, DEVICE
 
-    import torch
-    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+# ── Loaders ─────────────────────────────────────────────────────────────────
 
-    model_id = "Qwen/Qwen2.5-VL-3B-Instruct"
-    PROCESSOR = AutoProcessor.from_pretrained(
-        model_id,
-        min_pixels=64 * 28 * 28,
-        max_pixels=512 * 28 * 28,  # cap resolution for speed
+def _get_yolo():
+    global _YOLO
+    if _YOLO is not None:
+        return _YOLO
+    try:
+        from huggingface_hub import hf_hub_download
+        from ultralytics import YOLO
+        path = hf_hub_download(
+            repo_id="zyddnys/manga-image-translator",
+            filename="comictextdetector.pt",
+        )
+        _YOLO = YOLO(path)
+        print("✅ YOLO comic-text-detector loaded")
+    except Exception as exc:
+        print(f"⚠️  YOLO failed ({exc}) — using full-page OCR fallback")
+        _YOLO = "fallback"
+    return _YOLO
+
+
+def _get_ocr():
+    global _OCR_MODEL, _OCR_PROC, _DEVICE
+    if _OCR_MODEL is not None:
+        return _OCR_MODEL, _OCR_PROC, _DEVICE
+    import torch
+    from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+    _OCR_PROC = AutoProcessor.from_pretrained(
+        "Qwen/Qwen2.5-VL-7B-Instruct",
+        min_pixels=128 * 28 * 28,
+        max_pixels=512 * 28 * 28,
     )
-    MODEL = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-        model_id,
+    _OCR_MODEL = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        "Qwen/Qwen2.5-VL-7B-Instruct",
         torch_dtype=torch.float16,
         device_map="auto",
     ).eval()
-    DEVICE = next(MODEL.parameters()).device
-    return MODEL, PROCESSOR, DEVICE
+    _DEVICE = next(_OCR_MODEL.parameters()).device
+    print("✅ Qwen2.5-VL-7B loaded")
+    return _OCR_MODEL, _OCR_PROC, _DEVICE
 
 
-PROMPT = "Extract all text from this image. Return only the text, no explanations."
+# ── Detection helpers ────────────────────────────────────────────────────────
 
+def _yolo_detect(img: Image.Image, yolo) -> list[tuple[int, int, int, int]]:
+    """Return [(x1,y1,x2,y2)] sorted top-to-bottom, right-to-left."""
+    if yolo == "fallback":
+        return []
+    try:
+        arr = np.array(img)
+        results = yolo(arr, verbose=False, conf=0.25, iou=0.45)
+        boxes = []
+        for r in results:
+            for xyxy in r.boxes.xyxy.cpu().numpy():
+                x1, y1, x2, y2 = (int(v) for v in xyxy)
+                if (x2 - x1) > 15 and (y2 - y1) > 10:
+                    boxes.append((x1, y1, x2, y2))
+        # Sort: top-to-bottom, right-to-left (manhwa reading order)
+        boxes.sort(key=lambda b: (b[1] // 60, -b[0]))
+        return boxes
+    except Exception as exc:
+        print(f"YOLO detect error: {exc}")
+        return []
 
-@spaces.GPU(duration=60)
-def ocr_batch(images_b64_json: str) -> str:
-    """
-    Receive a JSON list of base64-encoded PNG images.
-    Returns a JSON list of extracted text strings (same order).
-    Processes all images in a single GPU forward pass for speed.
-    """
+
+def _safe_crop(img: Image.Image, box: tuple[int, int, int, int]) -> Image.Image | None:
+    w, h = img.size
+    x1, y1, x2, y2 = box
+    x1, y1 = max(0, x1 - 4), max(0, y1 - 4)
+    x2, y2 = min(w, x2 + 4), min(h, y2 + 4)
+    if x2 - x1 < 10 or y2 - y1 < 10:
+        return None
+    return img.crop((x1, y1, x2, y2))
+
+
+# ── OCR helper ───���───────────────────────────────────────────────────────────
+
+_BUBBLE_PROMPT = (
+    "Extract the text from this manga/manhwa speech bubble. "
+    "Output only the raw text, nothing else."
+)
+_PAGE_PROMPT = (
+    "This is a manga/manhwa page. Extract ALL text visible (dialogue, captions, SFX). "
+    "Output each text bubble or caption on its own line, in reading order. "
+    "Output only the text, no labels or explanations."
+)
+
+
+def _ocr_batch(images: list[Image.Image], model, proc, device,
+               is_full_page: bool = False) -> list[str]:
+    """Batch-OCR a list of PIL images. Returns one string per image."""
+    if not images:
+        return []
     import torch
     from qwen_vl_utils import process_vision_info
 
+    prompt = _PAGE_PROMPT if is_full_page else _BUBBLE_PROMPT
+    texts_in, imgs_flat = [], []
+
+    for img in images:
+        msgs = [{"role": "user", "content": [
+            {"type": "image", "image": img},
+            {"type": "text", "text": prompt},
+        ]}]
+        texts_in.append(proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
+        img_inputs, _ = process_vision_info(msgs)
+        imgs_flat.extend(img_inputs)
+
+    inputs = proc(
+        text=texts_in, images=imgs_flat,
+        padding=True, return_tensors="pt",
+    ).to(device)
+
+    with torch.no_grad():
+        gen = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=False,
+            pad_token_id=proc.tokenizer.eos_token_id,
+        )
+
+    trimmed = [o[len(i):] for i, o in zip(inputs["input_ids"], gen)]
+    return [t.strip() for t in proc.batch_decode(trimmed, skip_special_tokens=True,
+                                                  clean_up_tokenization_spaces=False)]
+
+
+# ── Main endpoint ────────────────────────────────────────────────────────────
+
+@spaces.GPU(duration=120)
+def process_pages(pages_b64_json: str) -> str:
+    """
+    Receive a JSON list of base64-encoded full-page images.
+    Returns a JSON list of text strings (one per page).
+    """
     try:
-        images_b64: list[str] = json.loads(images_b64_json)
+        pages_b64: list[str] = json.loads(pages_b64_json)
     except Exception as exc:
-        return json.dumps({"error": f"Invalid JSON: {exc}"})
+        return json.dumps({"error": f"Bad JSON: {exc}"})
 
-    if not images_b64:
+    if not pages_b64:
         return json.dumps([])
 
-    model, processor, device = load_model()
+    yolo = _get_yolo()
+    model, proc, device = _get_ocr()
 
-    # Decode all images
-    pil_images: list[Image.Image] = []
-    decode_errors: dict[int, str] = {}
-    for idx, b64 in enumerate(images_b64):
+    page_results: list[str] = []
+
+    for b64 in pages_b64:
         try:
-            img_bytes = base64.b64decode(b64)
-            pil_images.append(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
+            page_img = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
         except Exception as exc:
-            pil_images.append(None)
-            decode_errors[idx] = str(exc)
-
-    # Build per-image messages
-    texts_in: list[str] = []
-    image_inputs_all: list = []
-    valid_indices: list[int] = []
-
-    for idx, img in enumerate(pil_images):
-        if img is None:
+            page_results.append(f"[DECODE_ERROR: {exc}]")
             continue
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": img},
-                    {"type": "text", "text": PROMPT},
-                ],
-            }
-        ]
-        text_prompt = processor.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        image_inputs, _ = process_vision_info(messages)
-
-        texts_in.append(text_prompt)
-        image_inputs_all.extend(image_inputs)
-        valid_indices.append(idx)
-
-    if not texts_in:
-        results = [decode_errors.get(i, "[DECODE_ERROR]") for i in range(len(images_b64))]
-        return json.dumps(results, ensure_ascii=False)
 
-    # Single batched forward pass
-    try:
-        inputs = processor(
-            text=texts_in,
-            images=image_inputs_all,
-            padding=True,
-            return_tensors="pt",
-        ).to(device)
-
-        with torch.no_grad():
-            generated_ids = model.generate(
-                **inputs,
-                max_new_tokens=128,
-                do_sample=False,
-                pad_token_id=processor.tokenizer.eos_token_id,
-            )
-
-        # Trim prompt tokens
-        trimmed = [
-            out[len(inp):]
-            for inp, out in zip(inputs["input_ids"], generated_ids)
-        ]
-        decoded = processor.batch_decode(
-            trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-        )
+        try:
+            boxes = _yolo_detect(page_img, yolo)
+            crops = [c for b in boxes if (c := _safe_crop(page_img, b)) is not None]
+
+            if crops:
+                # Process in sub-batches of 8 to avoid OOM
+                SUBBATCH = 8
+                all_texts = []
+                for i in range(0, len(crops), SUBBATCH):
+                    all_texts.extend(_ocr_batch(crops[i:i+SUBBATCH], model, proc, device))
+                page_text = "\n".join(t for t in all_texts if t)
+            else:
+                page_text = ""
+
+            # Always fall back to full-page if result is empty
+            if not page_text.strip():
+                res = _ocr_batch([page_img], model, proc, device, is_full_page=True)
+                page_text = res[0] if res else ""
+
+            page_results.append(page_text)
 
-    except Exception as exc:
-        traceback.print_exc()
-        decoded = [f"[BATCH_ERROR: {exc}]"] * len(valid_indices)
+        except Exception as exc:
+            traceback.print_exc()
+            page_results.append(f"[PAGE_ERROR: {exc}]")
 
-    # Map back to original positions
-    results = [""] * len(images_b64)
-    for i, orig_idx in enumerate(valid_indices):
-        results[orig_idx] = decoded[i].strip() if i < len(decoded) else "[MISSING]"
-    for idx, err in decode_errors.items():
-        results[idx] = f"[DECODE_ERROR: {err}]"
+    return json.dumps(page_results, ensure_ascii=False)
 
-    return json.dumps(results, ensure_ascii=False)
 
+# ── Gradio UI ────────────────────────────────────────────────────────────────
 
 demo = gr.Interface(
-    fn=ocr_batch,
-    inputs=gr.Textbox(
-        label="Base64 Images (JSON array)",
-        placeholder='["<base64_1>", "<base64_2>", ...]',
-        lines=3,
-    ),
-    outputs=gr.Textbox(label="OCR Results (JSON array)", lines=5),
-    title="🦅 Eagle Eye — Manga OCR (Qwen2.5-VL-3B)",
-    description="Fast GPU OCR. Send base64 PNG images, get text back.",
+    fn=process_pages,
+    inputs=gr.Textbox(label="Pages JSON (base64 array)", lines=3),
+    outputs=gr.Textbox(label="Results JSON (text per page)", lines=10),
+    title="🦅 Eagle Eye — Manga OCR  (YOLO + Qwen2.5-VL-7B)",
+    description="Send full manga pages as base64 JSON → get text back.",
     flagging_mode="never",
 )