feat: hybrid ML allowlist rescue on WD needs-review

Co-authored-by: Cursor <cursoragent@cursor.com>

backend/app/api.py

@@ -58,6 +58,7 @@ from .taxonomy import (
     resolve_taxonomy_folder,
     taxonomy_folder_names,
 )
+from .hybrid_ml import merge_ml_allowlist_scores, should_run_hybrid_ml
 from .providers import probe_execution_providers
 from .storage import execute, fetch_all, fetch_one, from_json, to_json
 
@@ -100,6 +101,58 @@ def _assignment_noise_floor(confidence_threshold: float) -> float:
     return max(0.15, float(confidence_threshold) * 0.5)
 
 
+def _maybe_hybrid_ml_rescue(
+    result: _ImageInferenceResult,
+    matched_tags: set[str],
+    confidence_threshold: float,
+    experimental_media_enabled: bool,
+    tagger_model: str,
+    wd_general_threshold: float,
+    experimental_style_detector_enabled: bool,
+    hybrid_ml_on_review: bool,
+) -> _ImageInferenceResult:
+    """On WD needs_review, merge allowlisted ML scores and re-route."""
+    if not should_run_hybrid_ml(
+        enabled=hybrid_ml_on_review,
+        tagger_model=tagger_model,
+        needs_review=result.needs_review,
+        inference_failed=result.inference_failed,
+    ):
+        return result
+    try:
+        if experimental_media_enabled and is_experimental_media(result.image_path):
+            ml_scores = extract_scores_with_experimental_media(
+                result.image_path,
+                experimental_media_enabled,
+                tagger_model="ml_danbooru",
+                wd_general_threshold=wd_general_threshold,
+            )
+        else:
+            ml_scores = extract_scores(
+                result.image_path,
+                tagger_model="ml_danbooru",
+                wd_general_threshold=wd_general_threshold,
+            )
+        merged = merge_ml_allowlist_scores(result.scores, ml_scores)
+        rescued = _classify_from_scores(
+            result.image_path,
+            merged,
+            matched_tags,
+            confidence_threshold,
+            experimental_style_detector_enabled=experimental_style_detector_enabled,
+            hybrid_real_life=False,
+        )
+        note = "Hybrid ML allowlist rescue."
+        if rescued.reason:
+            rescued.reason = f"{rescued.reason} {note}"
+        elif rescued.needs_review != result.needs_review or rescued.primary_tag != result.primary_tag:
+            rescued.reason = note
+        return rescued
+    except Exception:
+        logger.exception("hybrid_ml_rescue_failed image=%s", result.image_path)
+        return result
+
+
 def _now_iso() -> str:
     return datetime.now(timezone.utc).isoformat()
 
@@ -634,6 +687,7 @@ def _infer_one_image(
     wd_general_threshold: float = 0.35,
     experimental_style_detector_enabled: bool = False,
     hybrid_real_life: bool = False,
+    hybrid_ml_on_review: bool = False,
 ) -> _ImageInferenceResult:
     if hybrid_real_life:
         return _infer_hybrid_one_image(
@@ -660,7 +714,7 @@ def _infer_one_image(
                 tagger_model=tagger_model,
                 wd_general_threshold=wd_general_threshold,
             )
-        return _classify_from_scores(
+        result = _classify_from_scores(
             image_path,
             scores,
             matched_tags,
@@ -668,6 +722,16 @@ def _infer_one_image(
             experimental_style_detector_enabled=experimental_style_detector_enabled,
             hybrid_real_life=False,
         )
+        return _maybe_hybrid_ml_rescue(
+            result,
+            matched_tags,
+            confidence_threshold,
+            experimental_media_enabled,
+            tagger_model,
+            wd_general_threshold,
+            experimental_style_detector_enabled,
+            hybrid_ml_on_review,
+        )
     except Exception as err:
         if _is_provider_related_error(err):
             logger.exception("inference_provider_failure image=%s", image_path)
@@ -695,6 +759,7 @@ def _infer_batch_with_fallback(
     wd_general_threshold: float = 0.35,
     experimental_style_detector_enabled: bool = False,
     hybrid_real_life: bool = False,
+    hybrid_ml_on_review: bool = False,
 ) -> tuple[list[_ImageInferenceResult], float, str]:
     if not image_paths:
         return [], 0.0, "none"
@@ -738,6 +803,7 @@ def _infer_batch_with_fallback(
                 wd_general_threshold,
                 experimental_style_detector_enabled,
                 False,
+                hybrid_ml_on_review,
             )
         ]
         elapsed_ms = (time.perf_counter() - start) * 1000.0
@@ -757,6 +823,7 @@ def _infer_batch_with_fallback(
                 wd_general_threshold,
                 experimental_style_detector_enabled,
                 False,
+                hybrid_ml_on_review,
             )
             for p in image_paths
         ]
@@ -773,13 +840,22 @@ def _infer_batch_with_fallback(
         if len(scores_by_image) != len(image_paths):
             raise RuntimeError("Batch inference result count mismatch")
         rows = [
-            _classify_from_scores(
-                image_path,
-                scores,
+            _maybe_hybrid_ml_rescue(
+                _classify_from_scores(
+                    image_path,
+                    scores,
+                    matched_tags,
+                    confidence_threshold,
+                    experimental_style_detector_enabled=experimental_style_detector_enabled,
+                    hybrid_real_life=False,
+                ),
                 matched_tags,
                 confidence_threshold,
-                experimental_style_detector_enabled=experimental_style_detector_enabled,
-                hybrid_real_life=False,
+                experimental_media_enabled,
+                tagger_model,
+                wd_general_threshold,
+                experimental_style_detector_enabled,
+                hybrid_ml_on_review,
             )
             for image_path, scores in zip(image_paths, scores_by_image)
         ]
@@ -802,6 +878,7 @@ def _infer_batch_with_fallback(
                 wd_general_threshold,
                 experimental_style_detector_enabled,
                 False,
+                hybrid_ml_on_review,
             )
             right_rows, right_ms, _ = _infer_batch_with_fallback(
                 image_paths[mid:],
@@ -813,6 +890,7 @@ def _infer_batch_with_fallback(
                 wd_general_threshold,
                 experimental_style_detector_enabled,
                 False,
+                hybrid_ml_on_review,
             )
             return left_rows + right_rows, left_ms + right_ms, "batch_fallback"
         row = _infer_one_image(
@@ -824,6 +902,7 @@ def _infer_batch_with_fallback(
             wd_general_threshold,
             experimental_style_detector_enabled,
             False,
+            hybrid_ml_on_review,
         )
         elapsed_ms = (time.perf_counter() - start) * 1000.0
         return [row], elapsed_ms, "single_fallback"
@@ -851,6 +930,7 @@ def _settings_from_db() -> AppSettings:
         experimental_style_detector_enabled=bool(
             row.get("experimental_style_detector_enabled", 0)
         ),
+        hybrid_ml_on_review=bool(row.get("hybrid_ml_on_review", 1)),
         selected_tags=selected_tags,
         max_inference_workers=int(row.get("max_inference_workers") or 2),
         inference_batch_size=int(row.get("inference_batch_size") or 4),
@@ -1026,6 +1106,7 @@ def _execute_run(
     wd_general_threshold: float = 0.35,
     experimental_style_detector_enabled: bool = False,
     real_life_filter: bool = False,
+    hybrid_ml_on_review: bool = False,
 ) -> None:
     try:
         # Mark running immediately so clients can cancel during provider probe / scan.
@@ -1174,6 +1255,7 @@ def _execute_run(
                         wd_general_threshold,
                         experimental_style_detector_enabled,
                         real_life_filter,
+                        False if real_life_filter else hybrid_ml_on_review,
                     )
                     pending[future] = next_batch
 
@@ -1384,6 +1466,7 @@ def _execute_reclassify(
     tagger_model: str = "wd_eva02_large",
     wd_general_threshold: float = 0.35,
     experimental_style_detector_enabled: bool = False,
+    hybrid_ml_on_review: bool = False,
 ) -> None:
     try:
         # Claim already set status=running; only refresh bookkeeping here.
@@ -1513,6 +1596,8 @@ def _execute_reclassify(
                         tagger_model,
                         wd_general_threshold,
                         experimental_style_detector_enabled,
+                        False,
+                        hybrid_ml_on_review,
                     )
                     pending[future] = next_batch
 
@@ -1671,7 +1756,7 @@ def save_settings(payload: SaveSettingsRequest) -> AppSettings:
             UPDATE settings
             SET root_repo = ?, categories_root = ?, confidence_threshold = ?,
                 default_migrate_mode = ?, scan_recursive = ?, experimental_media_enabled = ?,
-                experimental_style_detector_enabled = ?,
+                experimental_style_detector_enabled = ?, hybrid_ml_on_review = ?,
                 selected_tags_json = ?, max_inference_workers = ?, inference_batch_size = ?,
                 force_cpu_inference = ?, tagger_model = ?, wd_general_threshold = ?
             WHERE id = 1
@@ -1684,6 +1769,7 @@ def save_settings(payload: SaveSettingsRequest) -> AppSettings:
                 1 if payload.scan_recursive else 0,
                 1 if payload.experimental_media_enabled else 0,
                 1 if payload.experimental_style_detector_enabled else 0,
+                1 if payload.hybrid_ml_on_review else 0,
                 to_json(cleaned_tags),
                 int(payload.max_inference_workers),
                 int(payload.inference_batch_size),
@@ -1870,6 +1956,7 @@ def start_run(payload: StartRunRequest) -> StartRunResponse:
             "wd_general_threshold": current.wd_general_threshold,
             "experimental_style_detector_enabled": style_detector_enabled,
             "real_life_filter": real_life_filter,
+            "hybrid_ml_on_review": bool(current.hybrid_ml_on_review),
         },
         daemon=True,
     )
@@ -2010,6 +2097,7 @@ def reclassify_run(run_id: int, payload: ReclassifyRequest) -> ReclassifyRespons
             payload.tagger_model,
             current.wd_general_threshold,
             current.experimental_style_detector_enabled,
+            current.hybrid_ml_on_review,
         ),
         daemon=True,
     )

backend/app/hybrid_ml.py

@@ -0,0 +1,155 @@
+"""WD-primary + ML allowlist rescue on needs_review.
+
+Benchmark-driven: ML wins recall on character/act cues (loli, fellatio,
+tentacles, monster_girl, …) but has higher Voyeur soft-tag FPR. Never merge
+those soft tags from ML into WD scores.
+"""
+
+from __future__ import annotations
+
+# Tags where ML recall@0.35 beat WD and FP rates stayed low on the curated suite.
+HYBRID_ML_ALLOWLIST: frozenset[str] = frozenset(
+    {
+        # Character / age
+        "loli",
+        "shota",
+        # Oral
+        "fellatio",
+        "irrumatio",
+        "deepthroat",
+        "licking_penis",
+        "cooperative_fellatio",
+        "after_fellatio",
+        # Creampie / nakadashi evidence
+        "nakadashi",
+        "cum_in_pussy",
+        "internal_cumshot",
+        "after_vaginal",
+        # Tentacle act (not bare "tentacles")
+        "tentacle_sex",
+        "consensual_tentacles",
+        "tentacles_on_male",
+        "tentacle_pit",
+        "tentacles_under_clothes",
+        # Monster-girl catch-all identity
+        "monster_girl",
+        "monster_boy",
+        "slime_girl",
+        "slime_(creature)",
+        "lamia",
+        "harpy",
+        "scylla",
+        "spider_girl",
+        "arthropod_girl",
+        "dragon_girl",
+        "plant_girl",
+        "fish_girl",
+        "shark_girl",
+        "frog_girl",
+        "bird_girl",
+        "moth_girl",
+        "centaur",
+        "mermaid",
+        "female_goblin",
+        "minotaur",
+        "orc",
+        "cyclops",
+        "oni",
+        "traditional_youkai",
+        "demon_girl",
+        "vampire",
+        "werewolf",
+        "monsterification",
+        # Android catch-all identity
+        "android",
+        "robot_girl",
+        "humanoid_robot",
+        "cyborg",
+        "mecha_musume",
+        "mechanization",
+        # Other preferred folders with strong ML recall
+        "furry",
+        "pokemon_(creature)",
+        "bestiality",
+        "animal_penis",
+        "incest",
+        "twincest",
+        "netorare",
+        "cheating_(relationship)",
+        "impregnation",
+        "fertilization",
+        "ovum",
+    }
+)
+
+# Never trust ML for these even if someone expands the allowlist later.
+HYBRID_ML_DENYLIST: frozenset[str] = frozenset(
+    {
+        "cleavage",
+        "nude",
+        "completely_nude",
+        "pussy",
+        "ass",
+        "breasts",
+        "large_breasts",
+        "huge_breasts",
+        "medium_breasts",
+        "small_breasts",
+        "nipples",
+        "areolae",
+        "panties",
+        "underwear",
+        "bra",
+        "bikini",
+        "swimsuit",
+        "lingerie",
+        "upskirt",
+        "see-through",
+        "see_through",
+        "cameltoe",
+        "underboob",
+        "sideboob",
+        "topless",
+        "bottomless",
+        "no_panties",
+        "spread_legs",
+        "from_below",
+        "looking_at_viewer",
+    }
+)
+
+
+def should_run_hybrid_ml(
+    *,
+    enabled: bool,
+    tagger_model: str,
+    needs_review: bool,
+    inference_failed: bool,
+) -> bool:
+    if not enabled or not needs_review or inference_failed:
+        return False
+    return str(tagger_model or "").startswith("wd_")
+
+
+def merge_ml_allowlist_scores(
+    wd_scores: dict[str, float],
+    ml_scores: dict[str, float],
+    allowlist: frozenset[str] | set[str] | None = None,
+    denylist: frozenset[str] | set[str] | None = None,
+) -> dict[str, float]:
+    """Copy WD scores; for allowlisted tags take max(WD, ML). Deny Voyeur soft tags."""
+    allow = allowlist if allowlist is not None else HYBRID_ML_ALLOWLIST
+    deny = denylist if denylist is not None else HYBRID_ML_DENYLIST
+    merged = {str(k): float(v) for k, v in wd_scores.items()}
+    for tag, score in ml_scores.items():
+        key = str(tag)
+        if key in deny or key not in allow:
+            continue
+        try:
+            ml_val = float(score)
+        except (TypeError, ValueError):
+            continue
+        prev = float(merged.get(key, 0.0))
+        if ml_val > prev:
+            merged[key] = ml_val
+    return merged

backend/app/schemas.py

@@ -20,6 +20,9 @@ class AppSettings(BaseModel):
     experimental_media_enabled: bool = False
     # Dedicated real-vs-anime ONNX gate (imgutils / deepghs anime_real_cls).
     experimental_style_detector_enabled: bool = False
+    # When primary tagger is WD and an item needs review, re-run ML-Danbooru and
+    # merge only allowlisted high-recall tags (never Voyeur soft cues).
+    hybrid_ml_on_review: bool = True
     # Shuck3r-style persisted preferences (survive reload / restart).
     selected_tags: list[str] = Field(default_factory=list)
     # Shared ORT run lock; preprocess overlaps across workers. Prefer 2 on GPU.

backend/app/storage.py

@@ -82,6 +82,7 @@ def _ensure_settings_columns(conn: sqlite3.Connection) -> None:
         "scan_recursive": "INTEGER NOT NULL DEFAULT 1",
         "experimental_media_enabled": "INTEGER NOT NULL DEFAULT 0",
         "experimental_style_detector_enabled": "INTEGER NOT NULL DEFAULT 0",
+        "hybrid_ml_on_review": "INTEGER NOT NULL DEFAULT 1",
         "selected_tags_json": "TEXT NOT NULL DEFAULT '[]'",
         "max_inference_workers": "INTEGER NOT NULL DEFAULT 2",
         "inference_batch_size": "INTEGER NOT NULL DEFAULT 4",

backend/tests/test_api_run_progress.py

@@ -40,18 +40,21 @@ def test_settings_round_trip_includes_tagger_model(tmp_path: Path):
             "force_cpu_inference": False,
             "tagger_model": "wd_eva02_large",
             "wd_general_threshold": 0.4,
+            "hybrid_ml_on_review": False,
         }
         put_resp = client.put("/api/settings", json=payload)
         put_resp.raise_for_status()
         saved = put_resp.json()
         assert saved["tagger_model"] == "wd_eva02_large"
         assert saved["wd_general_threshold"] == 0.4
+        assert saved["hybrid_ml_on_review"] is False
         get_resp = client.get("/api/settings")
         get_resp.raise_for_status()
         loaded = get_resp.json()
         assert loaded["tagger_model"] == "wd_eva02_large"
         assert loaded["wd_general_threshold"] == 0.4
         assert loaded["max_inference_workers"] == 2
+        assert loaded["hybrid_ml_on_review"] is False
 
 
 def test_start_run_missing_root_returns_400(tmp_path: Path):

backend/tests/test_hybrid_ml.py

@@ -0,0 +1,125 @@
+"""Hybrid ML-on-review: allowlist merge must raise preferred tags without Voyeur noise."""
+
+from __future__ import annotations
+
+import pytest
+
+from app.hybrid_ml import (
+    HYBRID_ML_ALLOWLIST,
+    merge_ml_allowlist_scores,
+    should_run_hybrid_ml,
+)
+from app.taxonomy import choose_best_destination
+
+
+def test_merge_raises_allowlisted_ml_score_with_max() -> None:
+    wd = {"1girl": 0.99, "loli": 0.2}
+    ml = {"1girl": 0.5, "loli": 0.88}
+    merged = merge_ml_allowlist_scores(wd, ml)
+    assert merged["loli"] == pytest.approx(0.88)
+    assert merged["1girl"] == pytest.approx(0.99)
+
+
+def test_merge_ignores_voyeur_soft_tags_from_ml() -> None:
+    wd = {"1girl": 0.9, "cleavage": 0.4}
+    ml = {"cleavage": 0.99, "nude": 0.95, "pussy": 0.9, "fellatio": 0.8}
+    merged = merge_ml_allowlist_scores(wd, ml)
+    assert merged["cleavage"] == pytest.approx(0.4)
+    assert "nude" not in merged
+    assert "pussy" not in merged
+    assert merged["fellatio"] == pytest.approx(0.8)
+
+
+def test_merge_does_not_inject_non_allowlist_tags() -> None:
+    wd = {"solo": 0.9}
+    ml = {"blue_eyes": 0.99, "monster_girl": 0.7}
+    merged = merge_ml_allowlist_scores(wd, ml)
+    assert "blue_eyes" not in merged
+    assert merged["monster_girl"] == pytest.approx(0.7)
+    assert merged["solo"] == pytest.approx(0.9)
+
+
+def test_allowlist_covers_preferred_recall_tags() -> None:
+    for tag in (
+        "loli",
+        "shota",
+        "fellatio",
+        "irrumatio",
+        "cum_in_pussy",
+        "tentacle_sex",
+        "monster_girl",
+        "slime_girl",
+        "android",
+        "robot_girl",
+    ):
+        assert tag in HYBRID_ML_ALLOWLIST, tag
+
+
+def test_merged_scores_can_rescue_routing() -> None:
+    """WD miss + ML hit on allowlisted act should route after merge."""
+    selected = {"fellatio", "Voyeur", "sex"}
+    wd = {"sex": 0.4, "cleavage": 0.9}
+    ml = {"fellatio": 0.85, "cleavage": 0.99}
+    folder_before, _, _ = choose_best_destination(wd, selected)
+    assert folder_before != "fellatio"
+    merged = merge_ml_allowlist_scores(wd, ml)
+    folder_after, score, _ = choose_best_destination(merged, selected)
+    assert folder_after == "fellatio"
+    assert score >= 0.85
+
+
+def test_should_run_hybrid_ml_only_for_wd_review() -> None:
+    assert should_run_hybrid_ml(
+        enabled=True, tagger_model="wd_swinv2_v3", needs_review=True, inference_failed=False
+    )
+    assert not should_run_hybrid_ml(
+        enabled=True, tagger_model="ml_danbooru", needs_review=True, inference_failed=False
+    )
+    assert not should_run_hybrid_ml(
+        enabled=True, tagger_model="wd_swinv2_v3", needs_review=False, inference_failed=False
+    )
+    assert not should_run_hybrid_ml(
+        enabled=False, tagger_model="wd_swinv2_v3", needs_review=True, inference_failed=False
+    )
+    assert not should_run_hybrid_ml(
+        enabled=True, tagger_model="wd_swinv2_v3", needs_review=True, inference_failed=True
+    )
+
+
+def test_maybe_hybrid_ml_rescue_reclassifies(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None:
+    from pathlib import Path
+
+    from app.api import _ImageInferenceResult, _maybe_hybrid_ml_rescue
+
+    path = tmp_path / "x.jpg"
+    path.write_bytes(b"not-an-image")
+    base = _ImageInferenceResult(
+        image_path=path,
+        scores={"sex": 0.4, "cleavage": 0.9},
+        primary_tag=None,
+        primary_score=None,
+        secondary=[],
+        needs_review=True,
+        reason="No matching tags found among selected tags.",
+        inference_failed=False,
+    )
+
+    monkeypatch.setattr(
+        "app.api.extract_scores",
+        lambda *a, **k: {"fellatio": 0.9, "cleavage": 0.99},
+    )
+    out = _maybe_hybrid_ml_rescue(
+        base,
+        {"fellatio", "Voyeur", "sex"},
+        0.6,
+        False,
+        "wd_swinv2_v3",
+        0.35,
+        False,
+        True,
+    )
+    assert out.primary_tag == "fellatio"
+    assert out.needs_review is False
+    assert out.scores["fellatio"] == pytest.approx(0.9)
+    # ML Voyeur soft tag must not overwrite WD.
+    assert out.scores["cleavage"] == pytest.approx(0.9)

frontend/src/App.jsx

@@ -10,6 +10,7 @@ const DEFAULT_SETTINGS = {
   scan_recursive: true,
   experimental_media_enabled: false,
   experimental_style_detector_enabled: false,
+  hybrid_ml_on_review: true,
   selected_tags: [],
   max_inference_workers: 2,
   inference_batch_size: 8,
@@ -906,6 +907,20 @@ function App() {
                 Experimental: real vs anime style gate (dedicated ONNX →{" "}
                 <code>real_life</code>; uncertain → needs review)
               </label>
+              <label className="inline-check">
+                <input
+                  type="checkbox"
+                  checked={Boolean(settings.hybrid_ml_on_review)}
+                  onChange={(e) =>
+                    setSettings({
+                      ...settings,
+                      hybrid_ml_on_review: e.target.checked,
+                    })
+                  }
+                />
+                Hybrid: on needs-review (WD primary), re-run ML-Danbooru and merge
+                allowlisted high-recall tags only (skips Voyeur soft cues)
+              </label>
             </div>
           </fieldset>