Update app/services/classifier.py

app/services/classifier.py

@@ -18,18 +18,42 @@ def _parse_csv_env(raw_value: str, default_csv: str) -> list[str]:
 
 
 GAME_ZERO_SHOT_MODEL = os.getenv("GAME_ZERO_SHOT_MODEL", "google/siglip-so400m-patch14-384")
+# Game-focused labels for SigLIP text-image matching (optimized for game detection)
+# Multiple game types to capture diverse game screenshots, final decision delegated to Qwen LLM
 GAME_ZERO_SHOT_LABELS = _parse_csv_env(
     os.getenv("GAME_ZERO_SHOT_LABELS", ""),
-    "game screen,browser with work,code editor,video streaming,idle desktop",
+    "a video game with 3D graphics and player UI overlay,a casual mobile or browser game with colorful graphics,a multiplayer online game with HUD and game interface,a fast-paced action or FPS game with weapons and objectives,a puzzle or strategy game with game board or mechanics,a horror or dark-themed game with spooky visuals,a sports or racing game with athletic UI,an esports or competitive game interface,a role-playing game or RPG with character stats,a browser window with website and text,a desktop or idle screen",
 )
-GAME_ZERO_SHOT_GAME_LABEL = os.getenv("GAME_ZERO_SHOT_GAME_LABEL", "game screen").strip().lower()
-GAME_ZERO_SHOT_GAME_THRESHOLD = float(os.getenv("GAME_ZERO_SHOT_GAME_THRESHOLD", "0.78"))
-GAME_ZERO_SHOT_NOT_GAME_THRESHOLD = float(os.getenv("GAME_ZERO_SHOT_NOT_GAME_THRESHOLD", "0.72"))
+# Short label names for logic matching (primary: game labels, contrast: non-game)
+GAME_ZERO_SHOT_LABEL_NAMES = [
+    "game_3d",           # 0: game with 3D graphics
+    "game_casual",       # 1: casual/mobile games
+    "game_mmo",          # 2: multiplayer online
+    "game_action",       # 3: action/FPS
+    "game_puzzle",       # 4: puzzle/strategy
+    "game_horror",       # 5: horror/dark
+    "game_sports",       # 6: sports/racing
+    "game_esports",      # 7: esports/competitive
+    "game_rpg",          # 8: RPG
+    "browser_work",      # 9: contrast - browser
+    "idle_desktop",      # 10: contrast - idle
+]
+# Aggregate game score: max of all game labels
+GAME_ZERO_SHOT_GAME_LABEL_INDICES = [0, 1, 2, 3, 4, 5, 6, 7, 8]  # indices of game labels
+GAME_ZERO_SHOT_GAME_THRESHOLD = float(os.getenv("GAME_ZERO_SHOT_GAME_THRESHOLD", "0.35"))  # max of game labels
+GAME_ZERO_SHOT_NOT_GAME_THRESHOLD = float(os.getenv("GAME_ZERO_SHOT_NOT_GAME_THRESHOLD", "0.30"))  # contrast labels
+GAME_ZERO_SHOT_AMBIGUOUS_THRESHOLD = float(os.getenv("GAME_ZERO_SHOT_AMBIGUOUS_THRESHOLD", "0.20"))  # when scores too close
+
+# Qwen model for final decision gate
+QWEN_MODEL = os.getenv("QWEN_MODEL", "Qwen/Qwen2-0.5B-Instruct")
+QWEN_ENABLE = os.getenv("QWEN_ENABLE", "true").lower() == "true"
 
 _nsfw_runtime: dict[str, Any] | None = None
 _nsfw_error: str | None = None
 _game_zero_shot_runtime: dict[str, Any] | None = None
 _game_zero_shot_error: str | None = None
+_qwen_runtime: dict[str, Any] | None = None
+_qwen_error: str | None = None
 
 logger = logging.getLogger(__name__)
 
@@ -52,6 +76,13 @@ def warmup_game_classifiers() -> dict[str, str]:
     else:
         status["vision"] = f"unavailable:{_game_zero_shot_error or 'unknown'}"
 
+    if QWEN_ENABLE:
+        qwen_runtime = _load_qwen_runtime()
+        if qwen_runtime is not None:
+            status["decision"] = f"ready:{QWEN_MODEL}"
+        else:
+            status["decision"] = f"unavailable:{_qwen_error or 'unknown'}"
+
     logger.info("game-classifier warmup status=%s", status)
     return status
 
@@ -162,6 +193,94 @@ def _load_game_zero_shot_runtime() -> dict[str, Any] | None:
         return None
 
 
+def _load_qwen_runtime() -> dict[str, Any] | None:
+    """Load Qwen 0.5B model for game decision gate."""
+    global _qwen_runtime
+    global _qwen_error
+
+    if not QWEN_ENABLE:
+        return None
+    if _qwen_runtime is not None:
+        return _qwen_runtime
+    if _qwen_error is not None:
+        return None
+
+    try:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        tokenizer = AutoTokenizer.from_pretrained(QWEN_MODEL)
+        model = AutoModelForCausalLM.from_pretrained(
+            QWEN_MODEL,
+            torch_dtype="auto",
+            device_map="auto" if torch.cuda.is_available() else "cpu"
+        )
+        model.eval()
+
+        _qwen_runtime = {
+            "model": model,
+            "tokenizer": tokenizer,
+            "torch": torch,
+            "device": device,
+        }
+        logger.info("qwen runtime initialized model=%s device=%s", QWEN_MODEL, device)
+        return _qwen_runtime
+    except Exception as exc:
+        _qwen_error = str(exc)
+        logger.warning("qwen runtime init failed model=%s err=%s", QWEN_MODEL, _qwen_error)
+        return None
+
+
+def _classify_with_qwen_decision(file_path: str, siglip_breakdown: str, game_score: float) -> str:
+    """Use Qwen to make final game decision based on SigLIP breakdown."""
+    runtime = _load_qwen_runtime()
+    if runtime is None:
+        return "uncertain"  # fallback to manual review
+
+    model = runtime["model"]
+    tokenizer = runtime["tokenizer"]
+    torch = runtime["torch"]
+
+    try:
+        # Build prompt with SigLIP visual analysis
+        prompt = f"""You are a game detection AI. Analyze the following image classification scores and determine if this is a game screen.
+
+SigLIP Classification Results:
+{siglip_breakdown}
+
+Max Game Score: {game_score:.2%}
+
+Question: Is this a screenshot of a video game? Answer with only one word: "GAME" if it's a game, "NOT_GAME" if it's not.
+
+Answer:"""
+
+        # Tokenize and generate
+        inputs = tokenizer(prompt, return_tensors="pt").to(runtime["device"])
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=10,
+                temperature=0.3,
+                top_p=0.9,
+            )
+
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract answer
+        answer = response.strip().split('\n')[-1].strip().upper()
+        if "GAME" in answer and "NOT_GAME" not in answer:
+            return "game"
+        elif "NOT_GAME" in answer:
+            return "not_game"
+        else:
+            return "uncertain"
+
+    except Exception as exc:
+        logger.warning("qwen decision failed err=%s", repr(exc))
+        return "uncertain"
+
+
 def _classify_game_with_zero_shot(file_path: str, suspected_game: bool) -> dict[str, Any] | None:
     runtime = _load_game_zero_shot_runtime()
     if runtime is None:
@@ -174,12 +293,14 @@ def _classify_game_with_zero_shot(file_path: str, suspected_game: bool) -> dict[
     device = runtime["device"]
 
     try:
+        import torch.nn.functional as F
+        
         # Load and prepare image
         with Image.open(file_path) as img:
             image = img.convert("RGB")
 
-        # Prepare texts with the required template format
-        texts = [f"This is a photo of {label}." for label in GAME_ZERO_SHOT_LABELS]
+        # Use detailed visual descriptions as labels (no template needed)
+        texts = GAME_ZERO_SHOT_LABELS
         
         # Process inputs
         inputs = processor(text=texts, images=image, padding="max_length", return_tensors="pt")
@@ -190,7 +311,8 @@ def _classify_game_with_zero_shot(file_path: str, suspected_game: bool) -> dict[
             outputs = model(**inputs)
 
         logits_per_image = outputs.logits_per_image
-        probs = torch.sigmoid(logits_per_image)[0].cpu().tolist()
+        # Use softmax to normalize scores (sum=1.0, labels compete)
+        probs = F.softmax(logits_per_image[0], dim=0).cpu().tolist()
 
     except Exception as exc:
         logger.warning("game zero-shot inference failed err=%s", repr(exc))
@@ -199,31 +321,59 @@ def _classify_game_with_zero_shot(file_path: str, suspected_game: bool) -> dict[
     if not probs or len(probs) == 0:
         return None
 
+    # Calculate game score BEFORE sorting (using fixed indices of original probs)
+    game_score = 0.0
+    game_scores_breakdown = []
+    for idx in GAME_ZERO_SHOT_GAME_LABEL_INDICES:
+        if idx < len(probs):
+            score = float(probs[idx])
+            label = GAME_ZERO_SHOT_LABEL_NAMES[idx]
+            game_scores_breakdown.append(f"{label}={score:.3f}")
+            if score > game_score:
+                game_score = score
+
     # Normalize results to match expected format
     normalized: list[dict[str, Any]] = []
-    for label, score in zip(GAME_ZERO_SHOT_LABELS, probs):
-        normalized.append({"label": label.lower(), "score": float(score)})
+    for label_name, full_desc, score in zip(GAME_ZERO_SHOT_LABEL_NAMES, GAME_ZERO_SHOT_LABELS, probs):
+        normalized.append({"label": label_name, "description": full_desc, "score": float(score)})
 
     normalized.sort(key=lambda x: float(x["score"]), reverse=True)
     top = normalized[0]
-    game_score = 0.0
-    for row in normalized:
-        if row["label"] == GAME_ZERO_SHOT_GAME_LABEL:
-            game_score = float(row["score"])
-            break
 
-    visual_summary = "; ".join(f"{row['label']}={row['score']:.3f}" for row in normalized[:5])
+    visual_summary = "; ".join(f"{row['label']}={row['score']:.3f}" for row in normalized)
+    game_summary = " + ".join(game_scores_breakdown[:3])  # top 3 game scores
+
+    # Ambiguous case detection
+    # Check if multiple game labels have close scores
+    top_game_score = game_score
+    second_game_score = 0.0
+    for idx in GAME_ZERO_SHOT_GAME_LABEL_INDICES:
+        if idx < len(probs):
+            score = float(probs[idx])
+            if score < top_game_score and score > second_game_score:
+                second_game_score = score
+    
+    score_gap = top_game_score - second_game_score if top_game_score > 0 else 1.0
+    is_ambiguous = (
+        (game_score < GAME_ZERO_SHOT_GAME_THRESHOLD and 
+         score_gap < GAME_ZERO_SHOT_AMBIGUOUS_THRESHOLD) or
+        (game_score >= 0.15 and game_score < GAME_ZERO_SHOT_GAME_THRESHOLD and
+         float(top["score"]) < GAME_ZERO_SHOT_NOT_GAME_THRESHOLD)
+    )
 
+    # Clear game: high confidence
     if game_score >= GAME_ZERO_SHOT_GAME_THRESHOLD:
         return {
             "verdict": "game",
             "confidence": game_score,
-            "reason": "game_ui_terms",
+            "reason": "game_ui_detected",
             "source": f"zero-shot:{GAME_ZERO_SHOT_MODEL}",
             "visual_summary": visual_summary,
+            "game_breakdown": game_summary,
         }
 
-    if top["label"] in {"browser with work", "code editor", "idle desktop"} and float(top["score"]) >= GAME_ZERO_SHOT_NOT_GAME_THRESHOLD:
+    # Clear non-game: strong contrast signal
+    if top["label"] in {"browser_work", "idle_desktop"} and float(top["score"]) >= GAME_ZERO_SHOT_NOT_GAME_THRESHOLD:
         return {
             "verdict": "not_game",
             "confidence": float(top["score"]),
@@ -232,7 +382,38 @@ def _classify_game_with_zero_shot(file_path: str, suspected_game: bool) -> dict[
             "visual_summary": visual_summary,
         }
 
-    if suspected_game and game_score >= 0.45:
+    # Ambiguous case: close game scores, low overall scores
+    # Delegate to Qwen for final decision
+    if is_ambiguous:
+        logger.info(
+            "game-detect ambiguous case game_score=%.2f score_gap=%.3f top_label=%s top_score=%.2f, delegating to qwen",
+            game_score,
+            score_gap,
+            top["label"],
+            float(top["score"]),
+        )
+        
+        qwen_verdict = _classify_with_qwen_decision(file_path, visual_summary, game_score)
+        if qwen_verdict == "game":
+            return {
+                "verdict": "game",
+                "confidence": max(game_score, 0.45),
+                "reason": "qwen_game_decision",
+                "source": f"qwen:{QWEN_MODEL}+{GAME_ZERO_SHOT_MODEL}",
+                "visual_summary": visual_summary,
+                "game_breakdown": game_summary,
+            }
+        elif qwen_verdict == "not_game":
+            return {
+                "verdict": "not_game",
+                "confidence": max(float(top["score"]), 0.3),
+                "reason": "qwen_not_game_decision",
+                "source": f"qwen:{QWEN_MODEL}+{GAME_ZERO_SHOT_MODEL}",
+                "visual_summary": visual_summary,
+            }
+
+    # Final fallback if not ambiguous or Qwen uncertain
+    if suspected_game and game_score >= 0.25:
         return {
             "verdict": "uncertain",
             "confidence": max(game_score, 0.55),