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