app.py

import gradio as gr
import requests
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
import manga_ocr
import deepl
import os, textwrap, cv2, urllib.request

# ── Optional: download Wild Words font if not present ──────────────────────
FONT_PATH = "WildWords.ttf"
FONT_URL  = "https://github.com/zyddnys/manga-image-translator/raw/main/fonts/wildwords.ttf"

def ensure_font():
    if not os.path.exists(FONT_PATH):
        try:
            urllib.request.urlretrieve(FONT_URL, FONT_PATH)
        except Exception:
            pass  # fall back to PIL default

ensure_font()

# ── Init models ─────────────────────────────────────────────────────────────
_ocr   = None
_yolo  = None

def get_ocr():
    global _ocr
    if _ocr is None:
        _ocr = manga_ocr.MangaOcr()
    return _ocr

def get_yolo():
    global _yolo
    if _yolo is None:
        try:
            from ultralytics import YOLO
            # speech-bubble detection weights (community fine-tune)
            weights = "speech_bubble.pt"
            if not os.path.exists(weights):
                urllib.request.urlretrieve(
                    "https://huggingface.co/ogkalu/comic-speech-bubble-detector/resolve/main/speech_bubble.pt",
                    weights,
                )
            _yolo = YOLO(weights)
        except Exception:
            _yolo = None
    return _yolo

# ── Image loading ────────────────────────────────────────────────────────────
def load_image(url_or_file):
    if isinstance(url_or_file, str) and url_or_file.startswith("http"):
        r = requests.get(url_or_file, timeout=15)
        r.raise_for_status()
        return Image.open(BytesIO(r.content)).convert("RGB")
    return Image.fromarray(url_or_file).convert("RGB")

# ── Bubble detection: YOLO first, OpenCV fallback ───────────────────────────
def find_bubbles(img: Image.Image):
    yolo = get_yolo()
    if yolo is not None:
        results = yolo(np.array(img), verbose=False)[0]
        boxes = []
        for box in results.boxes.xyxy.cpu().numpy():
            x1, y1, x2, y2 = map(int, box)
            boxes.append((x1, y1, x2 - x1, y2 - y1))
        if boxes:
            return boxes

    # OpenCV fallback: flood-fill background from corners, leaving only
    # enclosed white regions (bubble interiors) as detectable blobs.
    arr = np.array(img.convert("L"))
    h_img, w_img = arr.shape
    _, bw = cv2.threshold(arr, 220, 255, cv2.THRESH_BINARY)

    bg = bw.copy()
    flood_mask = np.zeros((h_img + 2, w_img + 2), dtype=np.uint8)
    for seed in [(0, 0), (w_img - 1, 0), (0, h_img - 1), (w_img - 1, h_img - 1)]:
        if bg[seed[1], seed[0]] == 255:
            cv2.floodFill(bg, flood_mask, seed, 128)

    enclosed = np.where(bg == 255, 255, 0).astype(np.uint8)
    kernel = np.ones((3, 3), np.uint8)
    enclosed = cv2.morphologyEx(enclosed, cv2.MORPH_CLOSE, kernel, iterations=2)
    contours, _ = cv2.findContours(enclosed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    boxes = []
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        area = w * h
        if area < 2000 or area > 0.6 * h_img * w_img:
            continue
        if max(w, h) / max(min(w, h), 1) > 5:
            continue
        boxes.append((x, y, w, h))
    return boxes

# ── Inpaint: fill bubble interior white ─────────────────────────────────────
def inpaint_bubble(img: Image.Image, box) -> Image.Image:
    x, y, w, h = box
    draw = ImageDraw.Draw(img)
    draw.rectangle([x + 5, y + 5, x + w - 5, y + h - 5], fill="white")
    return img

# ── Render translated text, auto-sized to fit bubble ────────────────────────
def draw_text_in_bubble(img: Image.Image, box, text: str) -> Image.Image:
    x, y, w, h = box
    draw  = ImageDraw.Draw(img)
    pad   = 10
    max_w = w - pad * 2

    font_obj = None
    chosen_lines = [text]
    chosen_size  = 12

    for size in range(20, 7, -1):
        try:
            f = ImageFont.truetype(FONT_PATH, size) if os.path.exists(FONT_PATH) else ImageFont.load_default()
        except Exception:
            f = ImageFont.load_default()

        avg_char_w = size * 0.55
        chars_per_line = max(1, int(max_w / avg_char_w))
        lines = textwrap.wrap(text, width=chars_per_line) or [text]
        line_h = size + 5
        total_h = len(lines) * line_h

        if total_h <= h - pad * 2:
            font_obj     = f
            chosen_lines = lines
            chosen_size  = line_h
            break

    start_y = y + (h - len(chosen_lines) * chosen_size) // 2
    for i, line in enumerate(chosen_lines):
        bbox = draw.textbbox((0, 0), line, font=font_obj)
        tw   = bbox[2] - bbox[0]
        tx   = x + (w - tw) // 2
        ty   = start_y + i * chosen_size
        # thin black outline for readability
        for dx, dy in [(-1,-1),(1,-1),(-1,1),(1,1)]:
            draw.text((tx+dx, ty+dy), line, fill="white", font=font_obj)
        draw.text((tx, ty), line, fill="black", font=font_obj)
    return img

# ── Main pipeline ────────────────────────────────────────────────────────────
LANGS = {
    "English (US)": "EN-US",
    "English (UK)": "EN-GB",
    "Spanish":      "ES",
    "French":       "FR",
    "German":       "DE",
    "Portuguese":   "PT-BR",
    "Italian":      "IT",
    "Polish":       "PL",
}

def translate_manga(image_input, deepl_key: str, target_lang: str, detection_mode: str):
    if image_input is None:
        return None, "⚠️ Please provide an image URL or upload a file."

    try:
        img = load_image(image_input)
    except Exception as e:
        return None, f"❌ Could not load image: {e}"

    ocr     = get_ocr()
    bubbles = find_bubbles(img)

    if not bubbles:
        return img, "ℹ️ No speech bubbles detected. Try a cleaner scan."

    translator = deepl.Translator(deepl_key.strip()) if deepl_key.strip() else None
    lang_code  = LANGS.get(target_lang, "EN-US")
    log = [f"Found {len(bubbles)} bubble(s)\n"]

    for i, box in enumerate(bubbles):
        x, y, w, h = box
        crop = img.crop((x, y, x + w, y + h))
        try:
            jp_text = ocr(crop).strip()
        except Exception as e:
            log.append(f"  Bubble {i+1}: OCR error — {e}")
            continue

        if not jp_text:
            continue

        if translator:
            try:
                result  = translator.translate_text(jp_text, target_lang=lang_code)
                en_text = result.text
            except Exception as e:
                en_text = jp_text
                log.append(f"  Bubble {i+1}: translation error — {e}")
        else:
            en_text = f"[{jp_text}]"

        img = inpaint_bubble(img, box)
        img = draw_text_in_bubble(img, box, en_text)
        log.append(f'  Bubble {i+1}: “{jp_text}” → “{en_text}”')

    return img, "\n".join(log)

# ── Gradio UI ────────────────────────────────────────────────────────────────
with gr.Blocks(title="Manga Translator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📖 Manga Translator\nDrop an image URL **or** upload a file → translated speech bubbles.")

    with gr.Row():
        with gr.Column(scale=2):
            url_input  = gr.Textbox(label="Image URL", placeholder="https://…/page.jpg")
            file_input = gr.Image(label="…or upload image", type="numpy")
            deepl_key  = gr.Textbox(label="DeepL API Key (free tier OK)", type="password",
                                    placeholder="Leave blank to see OCR-only mode")
            lang_sel   = gr.Dropdown(list(LANGS.keys()), value="English (US)", label="Target language")
            det_mode   = gr.Radio(["Auto (YOLO → OpenCV fallback)", "OpenCV only"],
                                  value="Auto (YOLO → OpenCV fallback)", label="Detection mode")
            btn        = gr.Button("Translate", variant="primary")

        with gr.Column(scale=3):
            out_img = gr.Image(label="Translated page", type="pil")
            out_log = gr.Textbox(label="Log", lines=10)

    def run(url, file, key, lang, mode):
        src = url.strip() if url and url.strip() else file
        return translate_manga(src, key, lang, mode)

    btn.click(fn=run, inputs=[url_input, file_input, deepl_key, lang_sel, det_mode],
              outputs=[out_img, out_log])

    gr.Markdown("""
**Notes**
- Get a free DeepL key at [deepl.com/pro#developer](https://www.deepl.com/pro#developer) — 500k chars/month free.
- YOLO model downloads automatically on first run (~6 MB). OpenCV fallback works without it.
- Best results on clean B&W scans with white speech bubbles.
""")

demo.launch(server_name="0.0.0.0", server_port=7860)