app.py

"""
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

Input  : JSON list of base64-encoded full-page PNG images
Output : JSON list of text strings (one per page, bubbles separated by \\n)
"""

import base64
import io
import json
import traceback

import gradio as gr
import numpy as np
import spaces
from PIL import Image

# ── Global model cache (loaded once, kept alive between ZeroGPU calls) ─────

_YOLO = None        # ultralytics.YOLO  |  "fallback"
_OCR_MODEL = None
_OCR_PROC = None
_DEVICE = None


# ── Loaders ─────────────────────────────────────────────────────────────────

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,
    )
    _OCR_MODEL = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        "Qwen/Qwen2.5-VL-7B-Instruct",
        torch_dtype=torch.float16,
        device_map="auto",
    ).eval()
    _DEVICE = next(_OCR_MODEL.parameters()).device
    print("✅ Qwen2.5-VL-7B loaded")
    return _OCR_MODEL, _OCR_PROC, _DEVICE


# ── 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 []


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:
        pages_b64: list[str] = json.loads(pages_b64_json)
    except Exception as exc:
        return json.dumps({"error": f"Bad JSON: {exc}"})

    if not pages_b64:
        return json.dumps([])

    yolo = _get_yolo()
    model, proc, device = _get_ocr()

    page_results: list[str] = []

    for b64 in pages_b64:
        try:
            page_img = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
        except Exception as exc:
            page_results.append(f"[DECODE_ERROR: {exc}]")
            continue

        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()
            page_results.append(f"[PAGE_ERROR: {exc}]")

    return json.dumps(page_results, ensure_ascii=False)


# ── Gradio UI ────────────────────────────────────────────────────────────────

demo = gr.Interface(
    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",
)

if __name__ == "__main__":
    demo.launch()