fixed lookup file for tag-category

app.py

@@ -12,10 +12,8 @@ from __future__ import annotations
 
 import json
 import os
-import sqlite3
 import sys
 import tempfile
-import threading
 import time
 from pathlib import Path
 
@@ -39,6 +37,7 @@ def _import_hf_hub():
 
 HF_REPO = "realphongha/danbooru-tag-query"
 MODELS_DIR = "models"
+CATEGORY_JSON = "tag_category.json"
 
 IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
 IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
@@ -54,73 +53,51 @@ CATEGORY_MAP = {
 DEFAULT_TOP_K = None
 DEFAULT_MIN_SCORE = 0.2
 
-# ── TagCache (SQLite) ───────────────────────────────────────────────────────
+# ── category lookup ─────────────────────────────────────────────────────────
 
-_DB_PATH = Path("data/tag_cache.db")
+def load_category_map(checkpoint: str | Path) -> dict[str, int]:
+    """Load tag→category map. Tries HF hub download, then sidecar file.
 
+    Returns {tag: category_id} — all tags default to 0 (general).
+    """
+    ckpt = Path(checkpoint)
 
-class TagCache:
-    def __init__(self, db: str | Path = _DB_PATH):
-        self.db = Path(db)
-        self.db.parent.mkdir(parents=True, exist_ok=True)
-        self._conn = sqlite3.connect(str(self.db), check_same_thread=False)
-        self._lock = threading.Lock()
-        self._conn.execute(
-            """CREATE TABLE IF NOT EXISTS tag_cache (
-                name TEXT PRIMARY KEY,
-                category INTEGER,
-                wiki_body TEXT,
-                fetched_at REAL NOT NULL
-            )"""
-        )
-        self._conn.commit()
+    # 1) sidecar: model.onnx → tag_category.json beside it
+    if ckpt.suffix == ".onnx":
+        sidecar = ckpt.with_name(CATEGORY_JSON)
+        if sidecar.exists():
+            return json.loads(sidecar.read_text())
+
+    # 2) parent dir: dir/model.onnx → dir/tag_category.json
+    parent_sidecar = ckpt.parent / CATEGORY_JSON
+    if parent_sidecar.exists():
+        return json.loads(parent_sidecar.read_text())
+
+    # 3) HF hub: download alongside model variant
+    if ckpt.suffix == ".onnx":
+        # try to infer variant from path
+        parts = ckpt.parts
+        for i, p in enumerate(parts):
+            if p == MODELS_DIR and i + 2 < len(parts):
+                variant = parts[i + 1]
+                try:
+                    hf = _import_hf_hub()
+                    path = hf.hf_hub_download(
+                        repo_id=HF_REPO,
+                        filename=f"{MODELS_DIR}/{variant}/{CATEGORY_JSON}",
+                        repo_type="model",
+                    )
+                    return json.loads(Path(path).read_text())
+                except Exception:
+                    pass
+                break
 
-    def get(self, name: str) -> tuple[int | None, str | None] | None:
-        with self._lock:
-            cur = self._conn.execute(
-                "SELECT category, wiki_body FROM tag_cache WHERE name = ?", (name,)
-            )
-            return cur.fetchone()
-
-    def set(self, name: str, category: int | None, wiki_body: str | None):
-        with self._lock:
-            self._conn.execute(
-                """INSERT OR REPLACE INTO tag_cache (name, category, wiki_body, fetched_at)
-                   VALUES (?, ?, ?, ?)""",
-                (name, category, wiki_body, time.time()),
-            )
-            self._conn.commit()
-
-    def get_many(self, names: list[str]) -> dict[str, tuple[int | None, str | None]]:
-        if not names:
-            return {}
-        with self._lock:
-            placeholders = ",".join("?" for _ in names)
-            cur = self._conn.execute(
-                f"SELECT name, category, wiki_body FROM tag_cache "
-                f"WHERE name IN ({placeholders})", names,
-            )
-            return {row[0]: (row[1], row[2]) for row in cur}
-
-    def bulk_set(self, items: list[tuple[str, int | None, str | None]]):
-        with self._lock:
-            now = time.time()
-            self._conn.executemany(
-                """INSERT OR REPLACE INTO tag_cache (name, category, wiki_body, fetched_at)
-                   VALUES (?, ?, ?, ?)""",
-                [(name, cat, body, now) for name, cat, body in items],
-            )
-            self._conn.commit()
+    return {}
 
-    def clear(self):
-        with self._lock:
-            self._conn.execute("DELETE FROM tag_cache")
-            self._conn.commit()
 
-    def size(self) -> int:
-        with self._lock:
-            cur = self._conn.execute("SELECT COUNT(*) FROM tag_cache")
-            return cur.fetchone()[0]
+def get_category_name(cat_map: dict[str, int], tag: str) -> str:
+    cat_id = cat_map.get(tag, 0)
+    return CATEGORY_MAP.get(cat_id, "general")
 
 
 # ── image preprocessing ────────────────────────────────────────────────────
@@ -144,12 +121,6 @@ def preprocess(image: Image.Image, image_size: int = 448) -> np.ndarray:
 
 # ── sidecar loading ────────────────────────────────────────────────────────
 
-def _sidecar_path(checkpoint: Path, suffix: str) -> Path:
-    if checkpoint.suffix == ".onnx":
-        return checkpoint.with_name(checkpoint.stem + suffix)
-    return checkpoint / suffix.lstrip(".")
-
-
 def load_tag_to_id(checkpoint: str | Path) -> dict[str, int]:
     ckpt = Path(checkpoint)
     path = _sidecar_path(ckpt, ".tag_to_id.json")
@@ -170,6 +141,12 @@ def load_config(checkpoint: str | Path) -> dict:
     return json.loads(path.read_text())
 
 
+def _sidecar_path(checkpoint: Path, suffix: str) -> Path:
+    if checkpoint.suffix == ".onnx":
+        return checkpoint.with_name(checkpoint.stem + suffix)
+    return checkpoint / suffix.lstrip(".")
+
+
 # ── Predictor (ONNX) ──────────────────────────────────────────────────────
 
 class Predictor:
@@ -180,6 +157,7 @@ class Predictor:
         self.tag_to_id = load_tag_to_id(self.checkpoint)
         cfg = load_config(self.checkpoint)
         self.image_size = cfg.get("image_size", 448)
+        self.cat_map = load_category_map(self.checkpoint)
 
         providers = [
             ("CUDAExecutionProvider", {}),
@@ -203,6 +181,9 @@ class Predictor:
     def num_classes(self) -> int:
         return len(self.tag_to_id)
 
+    def category_name(self, tag: str) -> str:
+        return get_category_name(self.cat_map, tag)
+
 
 # ── model discovery & loading (HF hub) ─────────────────────────────────────
 
@@ -223,100 +204,24 @@ def discover_model_variants() -> list[str]:
         return []
 
 
-def resolve_model_path(variant: str) -> str:
-    onnx_path = f"{MODELS_DIR}/{variant}/model.onnx"
-    hf = _import_hf_hub()
-    return hf.hf_hub_download(
-        repo_id=HF_REPO,
-        filename=onnx_path,
-        repo_type="model",
-    )
-
-
-def resolve_sidecar_paths(variant: str) -> tuple[str, str]:
+def download_model_variant(variant: str) -> Path:
     hf = _import_hf_hub()
-    config_path = hf.hf_hub_download(
-        repo_id=HF_REPO,
-        filename=f"{MODELS_DIR}/{variant}/config.json",
-        repo_type="model",
-    )
-    tagmap_path = hf.hf_hub_download(
+    onnx_path = hf.hf_hub_download(
         repo_id=HF_REPO,
-        filename=f"{MODELS_DIR}/{variant}/tag_to_id.json",
+        filename=f"{MODELS_DIR}/{variant}/model.onnx",
         repo_type="model",
     )
-    return config_path, tagmap_path
-
-
-def download_model_variant(variant: str) -> Path:
-    onnx = Path(resolve_model_path(variant))
-    resolve_sidecar_paths(variant)
-    return onnx
-
-
-# ── tag enrichment (category lookup, no API) ───────────────────────────────
-
-_CACHE = TagCache()
-
-
-def prebuild_cache(tag_to_id: dict[str, int]):
-    """Pre-populate tag cache from HF dataset.
-
-    Every tag gets a category (defaults to 0 = general).
-    No lazy API calls needed.
-    """
-    try:
-        from datasets import load_dataset
-    except ImportError:
-        return 0, 0
-
-    tags = sorted(tag_to_id.keys(), key=lambda t: tag_to_id[t])
-
-    try:
-        ds = load_dataset("qdlabs/danbooru-tags", split="train")
-        cat_map = {row["name"]: row["category"] for row in ds}
-    except Exception:
-        return 0, 0
-
-    found = sum(1 for t in tags if t in cat_map)
-    todo = []
-    already = 0
-    for tag in tags:
-        cached = _CACHE.get(tag)
-        if cached is not None and cached[0] is not None:
-            already += 1
-            continue
-        cat = cat_map.get(tag, 0)   # 0 = general for unknown
-        existing_wiki = cached[1] if cached else None
-        todo.append((tag, cat, existing_wiki))
-
-    if todo:
-        _CACHE.bulk_set(todo)
-
-    return already + len(todo), found
-
-
-def enrich_tags(tags_scores: list[tuple[str, float]]) -> dict[str, dict]:
-    """Attach category info to each tag. No API calls."""
-    tags = [t for t, _ in tags_scores]
-    cached_map = _CACHE.get_many(tags)
-    result: dict[str, dict] = {}
-    for tag, score in tags_scores:
-        cached = cached_map.get(tag)
-        if cached is not None:
-            cat_id, _ = cached
-        else:
-            cat_id = None
-        result[tag] = {
-            "score": score,
-            "category": cat_id,
-            "category_name": CATEGORY_MAP.get(cat_id, "general"),
-        }
-    return result
-
-
-def format_tag(tag: str, use_underscore: bool) -> str:
-    return tag if use_underscore else tag.replace("_", " ")
+    # download sidecars (config, tag_to_id, tag_category) if they exist
+    for sidecar in ["config.json", "tag_to_id.json", CATEGORY_JSON]:
+        try:
+            hf.hf_hub_download(
+                repo_id=HF_REPO,
+                filename=f"{MODELS_DIR}/{variant}/{sidecar}",
+                repo_type="model",
+            )
+        except Exception:
+            pass  # optional — missing is fine
+    return Path(onnx_path)
 
 
 # ── Gradio UI ──────────────────────────────────────────────────────────────
@@ -343,7 +248,6 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
         gr.Markdown("# 🏷️ DanbooruTagQuery")
 
         with gr.Row():
-            # ── left: image + model ──
             with gr.Column(scale=1):
                 image_input = gr.Image(
                     label="Image",
@@ -368,7 +272,6 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
                 )
                 model_status = gr.Markdown("Ready")
 
-            # ── right: params + categories ──
             with gr.Column(scale=1):
                 top_k = gr.Number(
                     label="Top-K", value=DEFAULT_TOP_K, minimum=0, step=1
@@ -392,7 +295,6 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
                     value=["general"],
                 )
 
-        # ── outputs ──
         with gr.Tabs():
             with gr.TabItem("📋 Tag list"):
                 tag_table = gr.HTML(label="Tags")
@@ -401,12 +303,9 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
                     tag_string = gr.Textbox(label="Tags", lines=6, elem_id="csv-text")
                     copy_btn = gr.Button("📋", elem_id="copy-csv-btn")
 
-        # ── status row ──
         with gr.Row():
             status = gr.Markdown("Ready. Load an image and click **Analyze**.")
-            clear_cache_btn = gr.Button("🧹 Clear cache", size="sm", elem_id="clear-cache-btn")
 
-        # ── tag score query ──
         gr.Markdown("### 🔍 Tag score query")
         with gr.Row():
             tag_query = gr.Textbox(
@@ -458,7 +357,7 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
                     f"<td>{link}</td>"
                     f"<td>{display}</td>"
                     f"<td style='text-align:right'>{score:.4f}</td>"
-                    f"<td><code>{m['category_name'] or '?'}</code></td>"
+                    f"<td><code>{m['category_name']}</code></td>"
                     f"</tr>"
                 )
             table = (
@@ -503,7 +402,9 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
             t0 = time.time()
             all_logits = fn(pil)
             state["all_logits"] = all_logits
-            state["tag_metadata"] = enrich_tags(all_logits)
+            state["tag_metadata"] = enrich_tags(
+                all_logits, state["predictor"].cat_map
+            )
 
             table, csv = refresh_results(
                 top_k.value, min_score.value,
@@ -512,8 +413,7 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
             )
             elapsed = time.time() - t0
             n = len(state["tag_metadata"])
-            cached = _CACHE.size()
-            return table, csv, f"✅ {n} tags · {cached} cached · {elapsed:.2f}s"
+            return table, csv, f"✅ {n} tags · {elapsed:.2f}s"
 
         analyze_btn.click(
             fn=on_analyze,
@@ -540,12 +440,6 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
                      status, tag_query, tag_query_output],
         )
 
-        def on_clear_cache():
-            _CACHE.clear()
-            return "🧹 Cache cleared (0 entries)"
-
-        clear_cache_btn.click(fn=on_clear_cache, inputs=[], outputs=[status])
-
         for widget in [top_k, min_score, sort_by, use_underscore, categories]:
             widget.change(
                 fn=refresh_results,
@@ -566,7 +460,7 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
                 return "<i>No matching tags.</i>"
             rows = "".join(
                 f"<tr><td>{t}</td><td>{s:.4f}</td>"
-                f"<td><code>{meta[t]['category_name'] or '?'}</code></td></tr>"
+                f"<td><code>{meta[t]['category_name']}</code></td></tr>"
                 for t, s in matches
             )
             return (f"<table style='width:100%'>"
@@ -609,7 +503,6 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
                 state["predict_fn"] = new_predict_fn
                 state["all_logits"] = None
                 state["tag_metadata"] = None
-                _try_prebuild(predictor)
                 return f"✅ Switched to {variant} ({predictor.num_classes} tags)"
             except Exception as exc:
                 return f"❌ Failed to load model: {exc}"
@@ -623,16 +516,27 @@ def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
     return app
 
 
-# ── main ────────────────────────────────────────────────────────────────────
+# ── helpers ────────────────────────────────────────────────────────────────
+
+def enrich_tags(
+    tags_scores: list[tuple[str, float]], cat_map: dict[str, int]
+) -> dict[str, dict]:
+    result: dict[str, dict] = {}
+    for tag, score in tags_scores:
+        cat_id = cat_map.get(tag, 0)
+        result[tag] = {
+            "score": score,
+            "category": cat_id,
+            "category_name": CATEGORY_MAP.get(cat_id, "general"),
+        }
+    return result
 
-def _try_prebuild(predictor: Predictor):
-    try:
-        total, found = prebuild_cache(predictor.tag_to_id)
-        if total:
-            print(f"Cache prebuilt: {total} tags ({found} with category from HF dataset)")
-    except Exception as exc:
-        print(f"Cache prebuild skipped: {exc}")
 
+def format_tag(tag: str, use_underscore: bool) -> str:
+    return tag if use_underscore else tag.replace("_", " ")
+
+
+# ── main ────────────────────────────────────────────────────────────────────
 
 def main():
     model_arg = sys.argv[1] if len(sys.argv) > 1 else None
@@ -647,7 +551,6 @@ def main():
             sys.exit(1)
         print(f"Loading local model: {onnx}")
         predictor = Predictor(onnx)
-        _try_prebuild(predictor)
 
         def _predict(image: Image.Image) -> list[tuple[str, float]]:
             tensor = preprocess(image, predictor.image_size)
@@ -669,7 +572,6 @@ def main():
         onnx = onnx_files[0]
         print(f"Loading local model from MODEL_DIR: {onnx}")
         predictor = Predictor(onnx)
-        _try_prebuild(predictor)
 
         def _predict(image: Image.Image) -> list[tuple[str, float]]:
             tensor = preprocess(image, predictor.image_size)
@@ -691,7 +593,6 @@ def main():
             try:
                 onnx_path = download_model_variant(default)
                 predictor = Predictor(onnx_path)
-                _try_prebuild(predictor)
 
                 def _predict(image: Image.Image) -> list[tuple[str, float]]:
                     tensor = preprocess(image, predictor.image_size)

requirements.txt

@@ -3,5 +3,3 @@ gradio==6.20.0
 numpy==1.26.4
 pillow==12.3.0
 huggingface-hub==1.24.0
-datasets==5.0.0
-tqdm