File size: 2,604 Bytes
5487729
d249b39
2ebce2f
 
b301c0c
2ebce2f
 
d249b39
2ebce2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38e057b
2ebce2f
 
81fcb4f
 
38e057b
 
 
 
 
 
 
 
 
 
2ebce2f
 
fde7f25
5487729
2ebce2f
43099ab
5487729
2ebce2f
81fcb4f
 
 
 
 
 
 
 
 
 
 
 
 
5487729
2ebce2f
 
 
5487729
2ebce2f
 
 
81fcb4f
 
2ebce2f
5487729
 
2ebce2f
 
 
5487729
2ebce2f
38e057b
2ebce2f
 
 
3c6b13b
 
2ebce2f
 
 
5487729
 
 
b301c0c
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
import gradio as gr
import numpy as np
import cv2
import pandas as pd
import os
from pathlib import Path
from tempfile import NamedTemporaryFile

from src.pipeline_hf import process_single_image
from src.config_loader import load_config


# ================================
# Вспомогательная функция: подсветка сенсоров
# ================================
def draw_boxes(img_rgb, sensors):
    img = img_rgb.copy()

    for s in sensors:
        x, y, w, h = s["x"], s["y"], s["w"], s["h"]
        cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)

    return img


# ================================
# Основная функция демо
# ================================
def hf_process(file):

    cfg = load_config("configs/config.yaml")
    
    # Имя файла (если доступно)
    if hasattr(file, "name"):
        file_path = file.name
    else:
        file_path = file # Если пришла строка

    filename = Path(file_path).name

    img = cv2.imread(file_path)
    if img is None:
        raise ValueError("Не удалось прочитать изображение")

    # Запуск пайплайна
    title, sensors = process_single_image(img, color_ranges=cfg["colors"])

    # --- визуализация сенсоров ---
    boxed = draw_boxes(img, sensors)

    # --- DataFrame сенсоров ---
    df = pd.DataFrame([
        {
            "filename": filename,
            "title": title,
            "text": s["text"],
            "score": s["score"],
            "x": s["x"],
            "y": s["y"],
            "w": s["w"],
            "h": s["h"]
        }
        for s in sensors
    ])

    # --- Excel на скачивание ---
    tmp = NamedTemporaryFile(delete=False, suffix=".xlsx")
    df.to_excel(tmp.name, index=False)

    return (
        boxed,
        title,
        tmp.name,  # путь к скачиваемому файлу
        df
    )


# ================================
# Gradio UI
# ================================
demo = gr.Interface(
    fn=hf_process,
    inputs=gr.File(label="Upload image"),
    outputs=[
        gr.Image(label="Detected sensors"),
        gr.Textbox(label="Title"),
        gr.File(label="Download Excel"),
        gr.Dataframe(label="Recognized sensors")
    ],
    title="MNEMO OCR — HF Demo",
    description="Продовая демо-версия оцифровщика MNEMO."
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860))
    )