app.py

"""
Digital Inspector — Gradio Space
Runs on CPU (no GPU required for YOLOv8n inference)
"""

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
import pandas as pd
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

# ══════════════════════════════════════════════════════
# КОНФИГУРАЦИЯ
# ══════════════════════════════════════════════════════

MODEL_REPO = "AlihanSDev/digital-inspector"

MODELS = {
    "⭐ y8n_1024 — Best (fast + accurate, mAP=0.669)": ("y8n_1024", 1024),
    "y8s_1024  — High precision (precision=0.985)":    ("y8s_1024", 1024),
    "y8m_1024  — Medium (mAP=0.649)":                  ("y8m_1024", 1024),
    "y8l_1024  — Large (mAP=0.642)":                   ("y8l_1024", 1024),
    "y8s_640   — Fast, low resolution":                ("y8s_640",  640),
    "y8s_768   — Medium resolution":                   ("y8s_768",  768),
}

CLASS_NAMES_RU  = {0: "подпись", 1: "штамп",     2: "qr-код"}
CLASS_NAMES_EN  = {0: "signature", 1: "stamp",    2: "qr-code"}
CLASS_COLORS_RGB = {
    0: (255, 80,  80),
    1: (80,  180, 255),
    2: (80,  255, 130),
}

# ══════════════════════════════════════════════════════
# ПЕРЕВОДЫ
# ══════════════════════════════════════════════════════

T = {
    "RU": {
        "title":    "# 🔍 Digital Inspector",
        "subtitle": (
            "**Автоматическое обнаружение подписей, штампов и QR-кодов "
            "в деловых документах**\n\n"
            "Обучено на русско- и казахскоязычных документах.  \n"
            "Код: [GitHub](https://github.com/AlihanSDev/digital-inspector) | "
            "Модели: [HuggingFace](https://huggingface.co/AlihanSDev/digital-inspector)"
        ),
        "model_label":      "🤖 Модель",
        "conf_label":       "Порог уверенности",
        "iou_label":        "Порог IOU (NMS)",
        "lang_btn":         "🇬🇧 English",
        "tab_image":        "🖼️ Изображение (JPG / PNG)",
        "tab_pdf":          "📄 PDF документ",
        "img_input_label":  "Загрузи документ",
        "img_output_label": "Результат",
        "img_btn":          "🔍 Запустить детекцию",
        "pdf_input_label":  "Загрузи PDF",
        "page_label":       "Страница",
        "pdf_output_label": "Результат",
        "pdf_btn":          "🔍 Запустить детекцию",
        "summary_label":    "Итог",
        "table_label":      "Детекции",
        "no_image":   "Загрузи изображение",
        "no_pdf":     "Загрузи PDF файл",
        "found":      "✅ Найдено: ",
        "not_found":  "⚠️ Объекты не найдены. Попробуй снизить порог уверенности.",
        "page_of":    "Страница {cur}/{total} — ",
        "pdf_error":  "❌ Ошибка чтения PDF: ",
        "metrics_title": "📊 Метрики всех моделей",
        "about_title":   "ℹ️ Об исследовании",
        "metrics_md": """
## 📊 Результаты (Run 3 — консервативная аугментация)

| Модель | mAP@50 | mAP@50-95 | ms/img | Параметры |
|--------|--------|-----------|--------|-----------|
| **y8n_1024** ⭐ | **0.881** | **0.669** | **237** | 3.2M |
| y8m_1024 | 0.872 | 0.649 | 308 | 25.9M |
| y8l_1024 | 0.858 | 0.642 | 316 | 43.7M |
| y8s_1024 | 0.831 | 0.630 | 249 | 11.2M |
| y8s_640  | 0.823 | 0.615 | 225 | 11.2M |
| y8s_768  | 0.836 | 0.596 | 233 | 11.2M |

**Вывод:** y8n (3.2M параметров) обходит y8l (43.7M) по точности И скорости.
""",
        "about_md": """
## О проекте

Проведено **3 прогона** обучения:

| Прогон | Аугментация | mAP@50-95 |
|--------|-------------|-----------|
| Run 1 | Без аугментации | 0.650 |
| Run 2 | Mosaic + copy-paste | 0.252 ❌ |
| **Run 3** | **Консервативная** | **0.669** ✅ |

**Mosaic вреден для документов** — деградация в 2.5×.

### Классы
| Цвет | Класс | AP |
|------|-------|----|
| 🔴 | Подпись (signature) | 0.355 |
| 🔵 | Штамп (stamp) | 0.982 |
| 🟢 | QR-код | — |

### Ограничения
- Данные: деловые документы (рус/каз)
- Оптимальное разрешение: **1024px**
""",
    },
    "EN": {
        "title":    "# 🔍 Digital Inspector",
        "subtitle": (
            "**Automatic detection of signatures, stamps and QR codes "
            "in business documents**\n\n"
            "Trained on Russian and Kazakh language documents.  \n"
            "Code: [GitHub](https://github.com/AlihanSDev/digital-inspector) | "
            "Models: [HuggingFace](https://huggingface.co/AlihanSDev/digital-inspector)"
        ),
        "model_label":      "🤖 Model",
        "conf_label":       "Confidence threshold",
        "iou_label":        "IOU threshold (NMS)",
        "lang_btn":         "🇷🇺 Русский",
        "tab_image":        "🖼️ Image (JPG / PNG)",
        "tab_pdf":          "📄 PDF document",
        "img_input_label":  "Upload document",
        "img_output_label": "Result",
        "img_btn":          "🔍 Run detection",
        "pdf_input_label":  "Upload PDF",
        "page_label":       "Page",
        "pdf_output_label": "Result",
        "pdf_btn":          "🔍 Run detection",
        "summary_label":    "Summary",
        "table_label":      "Detections",
        "no_image":   "Upload an image",
        "no_pdf":     "Upload a PDF file",
        "found":      "✅ Found: ",
        "not_found":  "⚠️ Nothing found. Try lowering the confidence threshold.",
        "page_of":    "Page {cur}/{total} — ",
        "pdf_error":  "❌ PDF read error: ",
        "metrics_title": "📊 All model metrics",
        "about_title":   "ℹ️ About the research",
        "metrics_md": """
## 📊 Results (Run 3 — Conservative Augmentation)

| Model | mAP@50 | mAP@50-95 | ms/img | Params |
|-------|--------|-----------|--------|--------|
| **y8n_1024** ⭐ | **0.881** | **0.669** | **237** | 3.2M |
| y8m_1024 | 0.872 | 0.649 | 308 | 25.9M |
| y8l_1024 | 0.858 | 0.642 | 316 | 43.7M |
| y8s_1024 | 0.831 | 0.630 | 249 | 11.2M |
| y8s_640  | 0.823 | 0.615 | 225 | 11.2M |
| y8s_768  | 0.836 | 0.596 | 233 | 11.2M |

**Key finding:** y8n (3.2M params) beats y8l (43.7M) in both accuracy AND speed.
""",
        "about_md": """
## About the project

**3 training runs** with different augmentation strategies:

| Run | Augmentation | mAP@50-95 |
|-----|-------------|-----------|
| Run 1 | No augmentation | 0.650 |
| Run 2 | Mosaic + copy-paste | 0.252 ❌ |
| **Run 3** | **Conservative** | **0.669** ✅ |

**Mosaic augmentation is harmful for documents** — 2.5× degradation.

### Classes
| Color | Class | AP |
|-------|-------|----|
| 🔴 | Signature | 0.355 |
| 🔵 | Stamp | 0.982 |
| 🟢 | QR code | — |

### Limitations
- Data: Russian/Kazakh business documents
- Optimal input resolution: **1024px**
""",
    },
}

# ══════════════════════════════════════════════════════
# ЗАГРУЗКА МОДЕЛЕЙ (кэш)
# ══════════════════════════════════════════════════════

_model_cache: dict = {}


def load_model(model_folder: str) -> YOLO:
    if model_folder in _model_cache:
        return _model_cache[model_folder]
    print(f"📥 Loading {model_folder}...")
    pt_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=f"models/{model_folder}/best.pt",
    )
    # CPU inference — достаточно быстро для YOLOv8n
    model = YOLO(pt_path)
    _model_cache[model_folder] = model
    print(f"✅ {model_folder} ready!")
    return model


# ══════════════════════════════════════════════════════
# ИНФЕРЕНС (CPU)
# ══════════════════════════════════════════════════════

def run_detection(pil_image: Image.Image,
                  model_folder: str,
                  imgsz: int,
                  conf: float,
                  iou: float,
                  lang: str):
    model = load_model(model_folder)

    # device="" → автоматически CPU на HF Spaces
    results = model.predict(
        pil_image,
        imgsz=imgsz,
        conf=conf,
        iou=iou,
        device="cpu",
        verbose=False,
    )

    img_out = pil_image.copy().convert("RGB")
    draw = ImageDraw.Draw(img_out)
    class_names = CLASS_NAMES_RU if lang == "RU" else CLASS_NAMES_EN
    detections = []

    boxes = results[0].boxes
    if boxes is not None and len(boxes) > 0:
        for box in boxes:
            cls_id = int(box.cls[0])
            conf_v = float(box.conf[0])
            x1, y1, x2, y2 = [int(v) for v in box.xyxy[0]]
            color    = CLASS_COLORS_RGB.get(cls_id, (200, 200, 200))
            cls_name = class_names.get(cls_id, f"class_{cls_id}")

            draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
            label = f"{cls_name}  {conf_v:.2f}"
            lbl_w = len(label) * 8 + 4
            lbl_h = 22
            draw.rectangle(
                [x1, max(0, y1 - lbl_h), x1 + lbl_w, y1], fill=color)
            draw.text(
                (x1 + 2, max(0, y1 - lbl_h + 2)), label, fill=(0, 0, 0))

            detections.append({
                "class" if lang == "EN" else "класс": cls_name,
                "conf":  f"{conf_v:.3f}",
                "X1": x1, "Y1": y1, "X2": x2, "Y2": y2,
                "W" if lang == "EN" else "Ш": x2 - x1,
                "H" if lang == "EN" else "В": y2 - y1,
            })

    t = T[lang]
    if detections:
        counts: dict = {}
        for d in detections:
            key = d.get("class") or d.get("класс")
            counts[key] = counts.get(key, 0) + 1
        summary = t["found"] + ", ".join(f"{v} {k}" for k, v in counts.items())
    else:
        summary = t["not_found"]

    return np.array(img_out), pd.DataFrame(detections), summary


# ══════════════════════════════════════════════════════
# ОБРАБОТЧИКИ
# ══════════════════════════════════════════════════════

def handle_image(image_np, model_key, conf, iou, lang):
    t = T[lang]
    if image_np is None:
        return None, pd.DataFrame(), t["no_image"]
    model_folder, imgsz = MODELS[model_key]
    pil_img = Image.fromarray(image_np).convert("RGB")
    return run_detection(pil_img, model_folder, imgsz, conf, iou, lang)


def handle_pdf(pdf_file, model_key, conf, iou, page_num, lang):
    t = T[lang]
    if pdf_file is None:
        return None, pd.DataFrame(), t["no_pdf"]
    try:
        from pdf2image import convert_from_path
        pages = convert_from_path(pdf_file.name, dpi=200)
    except Exception as e:
        return None, pd.DataFrame(), t["pdf_error"] + str(e)

    total    = len(pages)
    page_num = max(1, min(int(page_num), total))
    pil_img  = pages[page_num - 1]

    model_folder, imgsz = MODELS[model_key]
    img_out, df, summary = run_detection(
        pil_img, model_folder, imgsz, conf, iou, lang)

    summary = t["page_of"].format(cur=page_num, total=total) + summary
    return img_out, df, summary


def handle_pdf_pages(pdf_file):
    if pdf_file is None:
        return gr.update(maximum=1, value=1)
    try:
        from pdf2image import convert_from_path
        pages = convert_from_path(pdf_file.name, dpi=72)
        return gr.update(maximum=len(pages), value=1)
    except Exception:
        return gr.update(maximum=1, value=1)


def switch_language(current_lang: str):
    new_lang = "EN" if current_lang == "RU" else "RU"
    t = T[new_lang]
    return (
        new_lang,
        gr.update(value=t["title"]),
        gr.update(value=t["subtitle"]),
        gr.update(label=t["model_label"]),
        gr.update(label=t["conf_label"]),
        gr.update(label=t["iou_label"]),
        gr.update(value=t["lang_btn"]),
        gr.update(label=t["img_input_label"]),
        gr.update(label=t["img_output_label"]),
        gr.update(value=t["img_btn"]),
        gr.update(label=t["pdf_input_label"]),
        gr.update(label=t["page_label"]),
        gr.update(label=t["pdf_output_label"]),
        gr.update(value=t["pdf_btn"]),
        gr.update(label=t["summary_label"]),
        gr.update(label=t["table_label"]),
        gr.update(label=t["summary_label"]),
        gr.update(label=t["table_label"]),
        gr.update(value=t["metrics_md"]),
        gr.update(value=t["about_md"]),
        gr.update(label=t["metrics_title"]),
        gr.update(label=t["about_title"]),
    )


# ══════════════════════════════════════════════════════
# GRADIO UI
# ══════════════════════════════════════════════════════

_t0 = T["RU"]

with gr.Blocks(title="Digital Inspector", theme=gr.themes.Soft()) as demo:

    lang_state = gr.State("RU")

    # ── Шапка ────────────────────────────
    with gr.Row():
        with gr.Column(scale=9):
            title_md    = gr.Markdown(_t0["title"])
            subtitle_md = gr.Markdown(_t0["subtitle"])
        with gr.Column(scale=1, min_width=130):
            lang_btn = gr.Button(_t0["lang_btn"], variant="secondary", size="sm")

    # ── Панель управления ─────────────────
    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=list(MODELS.keys()),
            value=list(MODELS.keys())[0],
            label=_t0["model_label"],
            scale=3,
        )
        conf_slider = gr.Slider(
            minimum=0.05, maximum=0.95, value=0.25, step=0.05,
            label=_t0["conf_label"], scale=2,
        )
        iou_slider = gr.Slider(
            minimum=0.1, maximum=0.95, value=0.45, step=0.05,
            label=_t0["iou_label"], scale=2,
        )

    # ── Вкладки ──────────────────────────
    with gr.Tabs():

        # Изображение
        with gr.TabItem(_t0["tab_image"]):
            with gr.Row():
                with gr.Column(scale=1):
                    img_input = gr.Image(
                        label=_t0["img_input_label"],
                        type="numpy", height=500)
                    img_btn = gr.Button(
                        _t0["img_btn"], variant="primary", size="lg")
                with gr.Column(scale=1):
                    img_output = gr.Image(
                        label=_t0["img_output_label"],
                        type="numpy", height=500)
            img_summary = gr.Textbox(
                label=_t0["summary_label"], interactive=False)
            img_table = gr.Dataframe(
                label=_t0["table_label"], interactive=False)

            # Только click — убираем дублирующий change
            img_btn.click(
                fn=handle_image,
                inputs=[img_input, model_dropdown,
                        conf_slider, iou_slider, lang_state],
                outputs=[img_output, img_table, img_summary],
                api_name="detect_image",
            )

        # PDF
        with gr.TabItem(_t0["tab_pdf"]):
            with gr.Row():
                with gr.Column(scale=1):
                    pdf_input = gr.File(
                        label=_t0["pdf_input_label"],
                        file_types=[".pdf"])
                    page_slider = gr.Slider(
                        minimum=1, maximum=1, value=1, step=1,
                        label=_t0["page_label"])
                    pdf_btn = gr.Button(
                        _t0["pdf_btn"], variant="primary", size="lg")
                with gr.Column(scale=1):
                    pdf_output = gr.Image(
                        label=_t0["pdf_output_label"],
                        type="numpy", height=500)
            pdf_summary = gr.Textbox(
                label=_t0["summary_label"], interactive=False)
            pdf_table = gr.Dataframe(
                label=_t0["table_label"], interactive=False)

            pdf_input.change(
                fn=handle_pdf_pages,
                inputs=[pdf_input],
                outputs=[page_slider],
                api_name="update_pages",
            )
            pdf_btn.click(
                fn=handle_pdf,
                inputs=[pdf_input, model_dropdown,
                        conf_slider, iou_slider, page_slider, lang_state],
                outputs=[pdf_output, pdf_table, pdf_summary],
                api_name="detect_pdf",
            )
            # Смена страницы — тоже только одна функция
            page_slider.change(
                fn=handle_pdf,
                inputs=[pdf_input, model_dropdown,
                        conf_slider, iou_slider, page_slider, lang_state],
                outputs=[pdf_output, pdf_table, pdf_summary],
                api_name="detect_pdf_page",
            )

    # ── Аккордеоны ───────────────────────
    with gr.Accordion(_t0["metrics_title"], open=False) as metrics_acc:
        metrics_md = gr.Markdown(_t0["metrics_md"])

    with gr.Accordion(_t0["about_title"], open=False) as about_acc:
        about_md = gr.Markdown(_t0["about_md"])

    # ── Переключение языка ────────────────
    lang_btn.click(
        fn=switch_language,
        inputs=[lang_state],
        outputs=[
            lang_state,
            title_md, subtitle_md,
            model_dropdown, conf_slider, iou_slider,
            lang_btn,
            img_input, img_output, img_btn,
            pdf_input, page_slider, pdf_output, pdf_btn,
            img_summary, img_table,
            pdf_summary, pdf_table,
            metrics_md, about_md,
            metrics_acc, about_acc,
        ],
        api_name="switch_lang",
    )

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