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