logo detection

requirements.txt

@@ -1,3 +1,3 @@
-altair
-pandas
-streamlit
+streamlit>=1.33.0
+huggingface_hub>=0.23.0
+pillow>=10.0.0

src/streamlit_app.py

@@ -1,40 +1,199 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import io
+import os
+from typing import Any, Dict, List, Optional
+
+from huggingface_hub import InferenceClient
+from PIL import Image, ImageDraw, ImageFont
 import streamlit as st
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+
+# App config
+st.set_page_config(
+    page_title="Logo Detector • Hugging Face",
+    page_icon="🔎",
+    layout="centered",
+)
+
+DEFAULT_MODEL_ID = "keremberke/yolov8m-logo-detector"
+
+
+def get_hf_token() -> Optional[str]:
+    """Prefer Streamlit secrets, then environment variable."""
+    token: Optional[str] = None
+    try:
+        # st.secrets behaves like a dict when available
+        token = st.secrets.get("HF_TOKEN")  # type: ignore[attr-defined]
+    except Exception:
+        token = None
+    if not token:
+        token = os.environ.get("HF_TOKEN")
+    return token
+
+
+@st.cache_resource(show_spinner=False)
+def get_client(model_id: str, token: Optional[str]) -> InferenceClient:
+    return InferenceClient(model=model_id, token=token)
+
+
+@st.cache_data(show_spinner=False, ttl=600)
+def run_detection(
+    model_id: str,
+    token: Optional[str],
+    image_bytes: bytes,
+) -> List[Dict[str, Any]]:
+    client = InferenceClient(model=model_id, token=token)
+    # The object_detection endpoint returns a list of dicts with keys:
+    # label, score, and box {xmin, ymin, xmax, ymax}
+    return client.object_detection(image=image_bytes)
+
+
+def draw_boxes(
+    image: Image.Image,
+    predictions: List[Dict[str, Any]],
+    threshold: float,
+) -> Image.Image:
+    annotated = image.copy()
+    draw = ImageDraw.Draw(annotated)
+    try:
+        font = ImageFont.load_default()
+    except Exception:
+        font = None  # Pillow will fallback
+
+    for pred in predictions:
+        score = float(pred.get("score", 0.0))
+        if score < threshold:
+            continue
+        label = str(pred.get("label", "logo"))
+        box = pred.get("box", {})
+        # Support alternative key names just in case
+        x0 = box.get("xmin", box.get("x_min", 0))
+        y0 = box.get("ymin", box.get("y_min", 0))
+        x1 = box.get("xmax", box.get("x_max", 0))
+        y1 = box.get("ymax", box.get("y_max", 0))
+        x0, y0, x1, y1 = float(x0), float(y0), float(x1), float(y1)
+
+        # Rectangle
+        draw.rectangle([(x0, y0), (x1, y1)], outline=(255, 0, 0), width=3)
+
+        # Label background
+        text = f"{label} {score:.2f}"
+        try:
+            # Compute text bounding box for background
+            text_bbox = draw.textbbox((int(x0), int(y0)), text, font=font)
+            tx0, ty0, tx1, _ = text_bbox
+        except Exception:
+            # Fallback: rough estimate for background width/height
+            tx0, ty0 = int(x0), int(y0) - 20
+            tx1 = int(x0) + 8 * len(text)
+        bg_top = min(ty0, y0)
+        bg_bottom = max(ty0, y0)
+        draw.rectangle(
+            [(tx0, bg_top - 2), (tx1, bg_bottom + 2)],
+            fill=(255, 0, 0),
+        )
+
+        # Text
+        draw.text(
+            (int(x0) + 2, int(y0) - 18),
+            text,
+            fill=(255, 255, 255),
+            font=font,
+        )
+
+    return annotated
+
+
+# Sidebar controls
+st.sidebar.header("⚙️ Настройки")
+model_id = st.sidebar.text_input(
+    "Hugging Face модель",
+    value=DEFAULT_MODEL_ID,
+    help="Например, YOLO модель для детекции логотипов",
+)
+threshold = st.sidebar.slider(
+    "Порог уверенности",
+    min_value=0.0,
+    max_value=1.0,
+    value=0.30,
+    step=0.01,
+)
+
+st.title("🔎 Поиск логотипов на изображении (Hugging Face • YOLO)")
+st.write(
+    "Загрузите изображение. Модель найдёт логотипы "
+    "и отрисует bounding boxes."
+)
+
+# Token hint
+hf_token = get_hf_token()
+if not hf_token:
+    st.info(
+        (
+            "Опционально укажите токен `HF_TOKEN` через `st.secrets` "
+            "или переменную окружения."
+        )
+    )
+
+uploaded = st.file_uploader(
+    "Выберите изображение",
+    type=["png", "jpg", "jpeg"],
+    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()
+
+    cols = st.columns(2)
+    with cols[0]:
+        st.image(image, caption="Оригинал", use_column_width=True)
+
+    with st.spinner("Детекция логотипов…"):
+        try:
+            predictions = run_detection(
+                model_id,
+                hf_token,
+                uploaded.getvalue(),
+            )
+        except Exception as exc:
+            st.error(f"Ошибка инференса: {exc}")
+            st.stop()
+
+    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_column_width=True,
+        )
+
+    # Stats and download
+    shown = sum(
+        1
+        for p in predictions
+        if float(p.get("score", 0.0)) >= threshold
+    )
+    total = len(predictions)
+    st.caption(
+        f"Показано боксов: {shown} из {total} "
+        f"(порог {threshold:.2f})"
+    )
+
+    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",
+    )