File size: 2,984 Bytes
6214b0d
 
4485c2a
6a9b444
 
4485c2a
a0709ea
4485c2a
6214b0d
 
 
4485c2a
 
 
6214b0d
 
 
 
ddb4219
6214b0d
ddb4219
b56bc5e
5f44bc6
6214b0d
ddb4219
6214b0d
 
 
 
4485c2a
6214b0d
 
 
ddb4219
6214b0d
5f44bc6
6214b0d
 
 
 
 
ddb4219
6214b0d
 
 
 
 
 
 
 
 
57f7dec
6214b0d
 
 
 
a0709ea
6214b0d
 
 
 
 
57f7dec
 
 
 
 
 
 
 
 
6214b0d
 
 
 
 
 
 
 
 
 
 
ddb4219
6214b0d
 
 
 
 
4485c2a
6214b0d
 
4485c2a
6214b0d
 
 
 
ddb4219
57f7dec
ddb4219
 
6214b0d
 
 
 
 
 
 
 
 
 
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
99
100
101
102
103
104
105
106
107
108
109
110
import io

from PIL import Image
import streamlit as st

import config
from utils import draw_boxes, run_detection


# App config
st.set_page_config(
    page_title=config.PAGE_TITLE,
    page_icon=config.PAGE_ICON,
    layout=config.LAYOUT,
)

# Sidebar controls
st.sidebar.header("⚙️ Настройки")
model_label = st.sidebar.selectbox(
    "Hugging Face модель",
    options=list(config.MODEL_CATALOG.keys()),
    index=0,
    help="Например, YOLO модель для детекции",
)
model_id = config.MODEL_CATALOG[model_label]
threshold = st.sidebar.slider(
    "Порог уверенности",
    min_value=0.0,
    max_value=1.0,
    value=float(config.DEFAULT_THRESHOLD),
    step=0.01,
)

st.title(config.PAGE_ICON + " " + config.PAGE_TITLE)
st.write(
    "Загрузите изображение. Модель найдёт объекты "
    "и отрисует bounding boxes."
)

uploaded = st.file_uploader(
    "Выберите изображение",
    type=config.UPLOADER_TYPES,
    accept_multiple_files=False,
)

if uploaded is not None:
    try:
        image = Image.open(uploaded).convert("RGB")
    except Exception as exc:
        st.error(f"Не удалось открыть изображение: {exc}")
        st.stop()
    
    with st.spinner("Детекция логотипов…"):
        try:
            predictions = run_detection(
                model_id,
                image,
            )
        except Exception as exc:
            st.error(f"Ошибка инференса: {exc}")
            st.stop()

    cols = st.columns(2)
    with cols[0]:
        st.image(
            image,
            caption="Оригинал",
            use_container_width=True,
        )


    if isinstance(predictions, dict) and predictions.get("error"):
        err_msg = predictions.get("error")
        st.error(f"Ошибка модели: {err_msg}")
        st.stop()

    annotated_image = draw_boxes(image, predictions, threshold)

    with cols[1]:
        st.image(
            annotated_image,
            caption="С найденными боксами",
            use_container_width=True,
        )

    # Stats and download
    shown = sum(
        1
        for p in predictions  # type: ignore[assignment]
        if float(p.get("score", 0.0)) >= threshold
    )
    total = len(predictions)  # type: ignore[arg-type]
    st.caption(
        f"Показано боксов: {shown} из {total} "
        f"(порог {threshold:.2f})"
    )
    predictions_str = "\n".join(
        [f"{p['label']}: {round(p['score'], 2)}" for p in predictions]
    )
    st.markdown(f"**{predictions_str}**")

    buf = io.BytesIO()
    annotated_image.save(buf, format="PNG")
    st.download_button(
        label="Скачать размеченное изображение",
        data=buf.getvalue(),
        file_name="detections.png",
        mime="image/png",
        type="primary",
    )