Update app/services/classifier.py

app/services/classifier.py

@@ -1,361 +1,387 @@
-from __future__ import annotations
-
-import json
-import os
-import re
-from pathlib import Path
-from typing import Any
-from urllib.parse import urlparse
-
-from app.core.config import settings
-
-MODELS_DIR = Path(__file__).resolve().parents[2] / "models"
-NSFW_CONFIG_PATH = MODELS_DIR / "config.json"
-NSFW_WEIGHTS_PATH = MODELS_DIR / "model.safetensors"
-NSFW_THRESHOLD = 0.75
-
-LLM_MODELS = [
-    os.getenv("GAME_LLM_MODEL", "Qwen/Qwen2.5-7B-Instruct"),
-    "microsoft/Phi-3.5-mini-instruct",
-]
-LLM_MAX_CHARS = 3000
-
-_nsfw_runtime: dict[str, Any] | None = None
-_nsfw_error: str | None = None
-_ocr_reader: Any | None = None
-_ocr_error: str | None = None
-
-GAME_KEYWORDS = {
-    "valorant",
-    "steam",
-    "roblox",
-    "league",
-    "dota",
-    "cs2",
-    "minecraft",
-    "epicgames",
-    "riot",
-    "crazygames",
-    "y8",
-    "miniclip",
-}
-
-SENSITIVE_KEYWORDS = {
-    "porn",
-    "sex",
-    "xxx",
-    "nsfw",
-    "adult",
-    "nude",
-    "erotic",
-}
-
-URL_REGEX = re.compile(r"(?:https?://)?(?:www\.)?[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}(?:/[^\s]*)?")
-
-GAME_PROMPT = """
-You are a strict classifier for parental-control screenshots.
-Input fields:
-- suspected_game_signal: boolean
-- extracted_urls: list of urls/domains from OCR
-- ocr_text: raw OCR text from screenshot
-Task:
-- Decide if screenshot likely indicates gaming/web-game activity.
-- Return only compact JSON with this schema:
-{
-  "verdict": "game" | "not_game" | "uncertain",
-  "confidence": 0.0-1.0,
-  "reason": "short reason"
-}
-Rules:
-- If clear game domain or game UI terms appear, lean game.
-- If evidence is weak or contradictory, return uncertain.
-- Never output markdown, prose, or extra keys.
-""".strip()
-
-
-def classify_screenshot(file_path: str, filename: str, suspected_game: bool) -> tuple[bool, float, str]:
-    game_result = classify_game_with_ocr_llm(file_path, filename, suspected_game)
-    verdict = game_result["verdict"]
-    confidence = float(game_result["confidence"])
-    reason = str(game_result["reason"])
-    if verdict == "game":
-        return True, confidence, reason
-    if verdict == "uncertain":
-        return True, max(confidence, 0.51), "uncertain-review"
-    return False, confidence, reason
-
-
-def classify_game_with_ocr_llm(file_path: str, filename: str, suspected_game: bool) -> dict[str, Any]:
-    ocr_text, urls = extract_ocr_text_and_urls(file_path)
-    llm = _classify_game_with_llm(ocr_text, urls, suspected_game)
-    if llm is not None:
-        return {
-            "verdict": llm["verdict"],
-            "confidence": llm["confidence"],
-            "reason": llm["reason"],
-            "ocr_text": ocr_text,
-            "urls": urls,
-            "source": llm.get("source", "llm"),
-        }
-
-    # Fallback heuristic path when model call is unavailable.
-    text_lower = ocr_text.lower()
-    keyword_hit = any(word in text_lower for word in GAME_KEYWORDS)
-    domain_hit = any(_domain_from_url(url) in GAME_KEYWORDS for url in urls)
-    if domain_hit or (suspected_game and keyword_hit):
-        return {
-            "verdict": "game",
-            "confidence": 0.78,
-            "reason": "ocr-keyword-heuristic",
-            "ocr_text": ocr_text,
-            "urls": urls,
-            "source": "heuristic",
-        }
-
-    if suspected_game:
-        return {
-            "verdict": "uncertain",
-            "confidence": 0.55,
-            "reason": "signal-without-clear-ocr",
-            "ocr_text": ocr_text,
-            "urls": urls,
-            "source": "heuristic",
-        }
-
-    return {
-        "verdict": "not_game",
-        "confidence": 0.2,
-        "reason": "no-game-evidence",
-        "ocr_text": ocr_text,
-        "urls": urls,
-        "source": "heuristic",
-    }
-
-
-def extract_ocr_text_and_urls(file_path: str) -> tuple[str, list[str]]:
-    reader = _load_ocr_reader()
-    if reader is None:
-        return "", []
-
-    try:
-        results = reader.readtext(file_path, detail=0, paragraph=True)
-        ocr_text = "\n".join(str(x) for x in results).strip()
-    except Exception:
-        return "", []
-
-    raw_urls = URL_REGEX.findall(ocr_text)
-    normalized = []
-    for value in raw_urls:
-        item = value.strip().rstrip(".,)")
-        if not item:
-            continue
-        if not item.startswith("http://") and not item.startswith("https://"):
-            item = f"https://{item}"
-        normalized.append(item.lower())
-
-    # Keep order stable while deduplicating.
-    urls: list[str] = []
-    seen: set[str] = set()
-    for item in normalized:
-        if item in seen:
-            continue
-        seen.add(item)
-        urls.append(item)
-
-    return ocr_text[:LLM_MAX_CHARS], urls
-
-
-def _load_ocr_reader() -> Any | None:
-    global _ocr_reader
-    global _ocr_error
-
-    if _ocr_reader is not None:
-        return _ocr_reader
-    if _ocr_error is not None:
-        return None
-
-    try:
-        import easyocr
-
-        _ocr_reader = easyocr.Reader(["en"], gpu=False)
-        return _ocr_reader
-    except Exception as exc:
-        _ocr_error = str(exc)
-        return None
-
-
-def _classify_game_with_llm(ocr_text: str, urls: list[str], suspected_game: bool) -> dict[str, Any] | None:
-    if not ocr_text and not urls:
-        return None
-
-    user_payload = {
-        "suspected_game_signal": suspected_game,
-        "extracted_urls": urls,
-        "ocr_text": ocr_text,
-    }
-
-    try:
-        from huggingface_hub import InferenceClient
-    except Exception:
-        return None
-
-    token = settings.hf_token
-    prompt_messages = [
-        {"role": "system", "content": GAME_PROMPT},
-        {"role": "user", "content": json.dumps(user_payload, ensure_ascii=True)},
-    ]
-
-    for model_name in LLM_MODELS:
-        try:
-            client = InferenceClient(model=model_name, token=token or None)
-            response = client.chat_completion(
-                messages=prompt_messages,
-                max_tokens=160,
-                temperature=0.1,
-                top_p=0.9,
-            )
-            content = ""
-            if response.choices:
-                content = response.choices[0].message.content or ""
-            parsed = _parse_llm_json(content)
-            if parsed is None:
-                continue
-            parsed["source"] = f"llm:{model_name}"
-            return parsed
-        except Exception:
-            continue
-
-    return None
-
-
-def _parse_llm_json(content: str) -> dict[str, Any] | None:
-    if not content:
-        return None
-
-    text = content.strip()
-    try:
-        data = json.loads(text)
-    except json.JSONDecodeError:
-        match = re.search(r"\{.*\}", text, re.DOTALL)
-        if not match:
-            return None
-        try:
-            data = json.loads(match.group(0))
-        except json.JSONDecodeError:
-            return None
-
-    verdict = str(data.get("verdict", "")).lower()
-    if verdict not in {"game", "not_game", "uncertain"}:
-        return None
-
-    confidence_raw = data.get("confidence", 0.5)
-    try:
-        confidence = float(confidence_raw)
-    except (TypeError, ValueError):
-        confidence = 0.5
-
-    confidence = max(0.0, min(1.0, confidence))
-    reason = str(data.get("reason", "llm-decision"))[:200]
-
-    return {"verdict": verdict, "confidence": confidence, "reason": reason}
-
-
-def _domain_from_url(url: str) -> str:
-    parsed = urlparse(url)
-    host = parsed.netloc or parsed.path
-    host = host.lower().replace("www.", "")
-    parts = host.split(".")
-    if not parts:
-        return host
-    return parts[0]
-
-
-def classify_sensitive_content(file_path: str, filename: str) -> tuple[bool, float, str]:
-    # First pass with image model inference; fallback to keyword signal only if unavailable.
-    model_result = _classify_sensitive_with_model(file_path)
-    if model_result is not None:
-        return model_result
-
-    name_text = f"{Path(file_path).name} {filename}".lower()
-    keyword_hit = any(word in name_text for word in SENSITIVE_KEYWORDS)
-    if keyword_hit:
-        return True, 0.65, "sensitive-keyword-fallback"
-    return False, 0.05, "no-sensitive-signal"
-
-
-def _load_nsfw_runtime() -> dict[str, Any] | None:
-    global _nsfw_runtime
-    global _nsfw_error
-
-    if _nsfw_runtime is not None:
-        return _nsfw_runtime
-    if _nsfw_error is not None:
-        return None
-
-    try:
-        import timm
-        import torch
-        from PIL import Image
-        from safetensors.torch import load_file
-
-        if not NSFW_CONFIG_PATH.exists() or not NSFW_WEIGHTS_PATH.exists():
-            _nsfw_error = "missing-local-model-files"
-            return None
-
-        config_data = json.loads(NSFW_CONFIG_PATH.read_text(encoding="utf-8"))
-        architecture = str(config_data.get("architecture", "vit_tiny_patch16_384"))
-        num_classes = int(config_data.get("num_classes", 2))
-        label_names = [str(x).lower() for x in config_data.get("label_names", ["nsfw", "sfw"])]
-        pretrained_cfg = config_data.get("pretrained_cfg", {})
-
-        model = timm.create_model(architecture, pretrained=False, num_classes=num_classes).eval()
-        state_dict = load_file(str(NSFW_WEIGHTS_PATH), device="cpu")
-        model.load_state_dict(state_dict, strict=False)
-
-        # Use local config for preprocessing so inference does not depend on remote metadata.
-        model.pretrained_cfg = {**getattr(model, "pretrained_cfg", {}), **pretrained_cfg, "label_names": label_names}
-
-        data_config = timm.data.resolve_model_data_config(model)
-        transforms = timm.data.create_transform(**data_config, is_training=False)
-
-        _nsfw_runtime = {
-            "torch": torch,
-            "Image": Image,
-            "model": model,
-            "transforms": transforms,
-            "label_names": [str(x).lower() for x in label_names],
-        }
-        return _nsfw_runtime
-    except Exception as exc:
-        _nsfw_error = str(exc)
-        return None
-
-
-def _classify_sensitive_with_model(file_path: str) -> tuple[bool, float, str] | None:
-    runtime = _load_nsfw_runtime()
-    if runtime is None:
-        return None
-
-    torch = runtime["torch"]
-    Image = runtime["Image"]
-    model = runtime["model"]
-    transforms = runtime["transforms"]
-    label_names = runtime["label_names"]
-
-    with Image.open(file_path) as img:
-        img = img.convert("RGB")
-        with torch.no_grad():
-            output = model(transforms(img).unsqueeze(0)).softmax(dim=-1).cpu()[0]
-
-    scores = [float(x) for x in output.tolist()]
-    nsfw_score = _extract_nsfw_score(scores, label_names)
-    return (nsfw_score >= NSFW_THRESHOLD, nsfw_score, "timm-marqo-nsfw")
-
-
-def _extract_nsfw_score(scores: list[float], labels: list[str]) -> float:
-    for idx, label in enumerate(labels):
-        if "nsfw" in label:
-            return scores[idx]
-    if len(scores) >= 2:
-        return scores[1]
-    return scores[0] if scores else 0.0
+from __future__ import annotations
+
+import json
+import os
+import re
+from pathlib import Path
+from typing import Any
+from urllib.parse import urlparse
+
+from app.core.config import settings
+
+MODELS_DIR = Path(__file__).resolve().parents[2] / "models"
+NSFW_CONFIG_PATH = MODELS_DIR / "config.json"
+NSFW_WEIGHTS_PATH = MODELS_DIR / "model.safetensors"
+NSFW_THRESHOLD = 0.75
+
+LLM_MODELS = [
+    os.getenv("GAME_LLM_MODEL", "Qwen/Qwen2.5-7B-Instruct"),
+    "microsoft/Phi-3.5-mini-instruct",
+]
+LLM_MAX_CHARS = 3000
+
+_nsfw_runtime: dict[str, Any] | None = None
+_nsfw_error: str | None = None
+_ocr_reader: Any | None = None
+_ocr_error: str | None = None
+
+GAME_KEYWORDS = {
+    "valorant",
+    "steam",
+    "roblox",
+    "league",
+    "dota",
+    "cs2",
+    "minecraft",
+    "epicgames",
+    "riot",
+    "crazygames",
+    "y8",
+    "miniclip",
+    "poki",
+    "friv",
+    "game",
+    "games",
+    "doodle",
+}
+
+SENSITIVE_KEYWORDS = {
+    "porn",
+    "sex",
+    "xxx",
+    "nsfw",
+    "adult",
+    "nude",
+    "erotic",
+}
+
+URL_REGEX = re.compile(r"(?:https?://)?(?:www\.)?[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}(?:/[^\s]*)?")
+
+GAME_PROMPT = """
+You are a strict classifier for parental-control screenshots.
+Input fields:
+- suspected_game_signal: boolean
+- extracted_urls: list of urls/domains from OCR
+- ocr_text: raw OCR text from screenshot
+Task:
+- Decide if screenshot likely indicates gaming/web-game activity.
+- Return only compact JSON with this schema:
+{
+  "verdict": "game" | "not_game" | "uncertain",
+  "confidence": 0.0-1.0,
+  "reason": "short reason"
+}
+Rules:
+- If clear game domain or game UI terms appear, lean game.
+- If evidence is weak or contradictory, return uncertain.
+- Never output markdown, prose, or extra keys.
+""".strip()
+
+
+def classify_screenshot(file_path: str, filename: str, suspected_game: bool) -> tuple[bool, float, str]:
+    game_result = classify_game_with_ocr_llm(file_path, filename, suspected_game)
+    verdict = game_result["verdict"]
+    confidence = float(game_result["confidence"])
+    reason = str(game_result["reason"])
+    if verdict == "game":
+        return True, confidence, reason
+    if verdict == "uncertain":
+        return True, max(confidence, 0.51), "uncertain-review"
+    return False, confidence, reason
+
+
+def classify_game_with_ocr_llm(file_path: str, filename: str, suspected_game: bool) -> dict[str, Any]:
+    ocr_text, urls = extract_ocr_text_and_urls(file_path)
+    llm = _classify_game_with_llm(ocr_text, urls, suspected_game)
+    if llm is not None:
+        return {
+            "verdict": llm["verdict"],
+            "confidence": llm["confidence"],
+            "reason": llm["reason"],
+            "ocr_text": ocr_text,
+            "urls": urls,
+            "source": llm.get("source", "llm"),
+        }
+
+    # Fallback heuristic path when model call is unavailable.
+    text_lower = ocr_text.lower()
+    keyword_hit = any(word in text_lower for word in GAME_KEYWORDS)
+    domain_hit = any(_is_game_like_domain(url) for url in urls)
+    if domain_hit or (suspected_game and keyword_hit):
+        return {
+            "verdict": "game",
+            "confidence": 0.78,
+            "reason": "ocr-keyword-heuristic",
+            "ocr_text": ocr_text,
+            "urls": urls,
+            "source": "heuristic",
+        }
+
+    if suspected_game:
+        return {
+            "verdict": "uncertain",
+            "confidence": 0.55,
+            "reason": "signal-without-clear-ocr",
+            "ocr_text": ocr_text,
+            "urls": urls,
+            "source": "heuristic",
+        }
+
+    return {
+        "verdict": "not_game",
+        "confidence": 0.2,
+        "reason": "no-game-evidence",
+        "ocr_text": ocr_text,
+        "urls": urls,
+        "source": "heuristic",
+    }
+
+
+def extract_ocr_text_and_urls(file_path: str) -> tuple[str, list[str]]:
+    reader = _load_ocr_reader()
+    if reader is None:
+        return "", []
+
+    try:
+        results = reader.readtext(file_path, detail=0, paragraph=True)
+        ocr_text = "\n".join(str(x) for x in results).strip()
+    except Exception:
+        return "", []
+
+    raw_urls = URL_REGEX.findall(ocr_text)
+    normalized = []
+    for value in raw_urls:
+        item = value.strip().rstrip(".,)")
+        if not item:
+            continue
+        if not item.startswith("http://") and not item.startswith("https://"):
+            item = f"https://{item}"
+        normalized.append(item.lower())
+
+    # Keep order stable while deduplicating.
+    urls: list[str] = []
+    seen: set[str] = set()
+    for item in normalized:
+        if item in seen:
+            continue
+        seen.add(item)
+        urls.append(item)
+
+    return ocr_text[:LLM_MAX_CHARS], urls
+
+
+def _load_ocr_reader() -> Any | None:
+    global _ocr_reader
+    global _ocr_error
+
+    if _ocr_reader is not None:
+        return _ocr_reader
+    if _ocr_error is not None:
+        return None
+
+    try:
+        import easyocr
+
+        _ocr_reader = easyocr.Reader(["en"], gpu=False)
+        return _ocr_reader
+    except Exception as exc:
+        _ocr_error = str(exc)
+        return None
+
+
+def _classify_game_with_llm(ocr_text: str, urls: list[str], suspected_game: bool) -> dict[str, Any] | None:
+    if not ocr_text and not urls:
+        return None
+
+    user_payload = {
+        "suspected_game_signal": suspected_game,
+        "extracted_urls": urls,
+        "ocr_text": ocr_text,
+    }
+
+    try:
+        from huggingface_hub import InferenceClient
+    except Exception:
+        return None
+
+    token = settings.hf_token
+    prompt_messages = [
+        {"role": "system", "content": GAME_PROMPT},
+        {"role": "user", "content": json.dumps(user_payload, ensure_ascii=True)},
+    ]
+
+    for model_name in LLM_MODELS:
+        try:
+            client = InferenceClient(model=model_name, token=token or None)
+            response = client.chat_completion(
+                messages=prompt_messages,
+                max_tokens=160,
+                temperature=0.1,
+                top_p=0.9,
+            )
+            content = ""
+            if response.choices:
+                content = response.choices[0].message.content or ""
+            parsed = _parse_llm_json(content)
+            if parsed is None:
+                continue
+            parsed["source"] = f"llm:{model_name}"
+            return parsed
+        except Exception:
+            continue
+
+    return None
+
+
+def _parse_llm_json(content: str) -> dict[str, Any] | None:
+    if not content:
+        return None
+
+    text = content.strip()
+    try:
+        data = json.loads(text)
+    except json.JSONDecodeError:
+        match = re.search(r"\{.*\}", text, re.DOTALL)
+        if not match:
+            return None
+        try:
+            data = json.loads(match.group(0))
+        except json.JSONDecodeError:
+            return None
+
+    verdict = str(data.get("verdict", "")).lower()
+    if verdict not in {"game", "not_game", "uncertain"}:
+        return None
+
+    confidence_raw = data.get("confidence", 0.5)
+    try:
+        confidence = float(confidence_raw)
+    except (TypeError, ValueError):
+        confidence = 0.5
+
+    confidence = max(0.0, min(1.0, confidence))
+    reason = str(data.get("reason", "llm-decision"))[:200]
+
+    return {"verdict": verdict, "confidence": confidence, "reason": reason}
+
+
+def _domain_from_url(url: str) -> str:
+    parsed = urlparse(url)
+    host = parsed.netloc or parsed.path
+    host = host.lower().replace("www.", "")
+    parts = host.split(".")
+    if not parts:
+        return host
+    return parts[0]
+
+
+def _is_game_like_domain(url: str) -> bool:
+    parsed = urlparse(url)
+    host = (parsed.netloc or parsed.path).lower().replace("www.", "")
+    first_label = _domain_from_url(url)
+
+    if first_label in GAME_KEYWORDS:
+        return True
+
+    # Catch common browser game hosts such as *.games, *game*, play-* styles.
+    host_tokens = re.split(r"[^a-z0-9]+", host)
+    for token in host_tokens:
+        if not token:
+            continue
+        if token in GAME_KEYWORDS:
+            return True
+        if "game" in token:
+            return True
+
+    return host.endswith(".games")
+
+
+def classify_sensitive_content(file_path: str, filename: str) -> tuple[bool, float, str]:
+    # First pass with image model inference; fallback to keyword signal only if unavailable.
+    model_result = _classify_sensitive_with_model(file_path)
+    if model_result is not None:
+        return model_result
+
+    name_text = f"{Path(file_path).name} {filename}".lower()
+    keyword_hit = any(word in name_text for word in SENSITIVE_KEYWORDS)
+    if keyword_hit:
+        return True, 0.65, "sensitive-keyword-fallback"
+    return False, 0.05, "no-sensitive-signal"
+
+
+def _load_nsfw_runtime() -> dict[str, Any] | None:
+    global _nsfw_runtime
+    global _nsfw_error
+
+    if _nsfw_runtime is not None:
+        return _nsfw_runtime
+    if _nsfw_error is not None:
+        return None
+
+    try:
+        import timm
+        import torch
+        from PIL import Image
+        from safetensors.torch import load_file
+
+        if not NSFW_CONFIG_PATH.exists() or not NSFW_WEIGHTS_PATH.exists():
+            _nsfw_error = "missing-local-model-files"
+            return None
+
+        config_data = json.loads(NSFW_CONFIG_PATH.read_text(encoding="utf-8"))
+        architecture = str(config_data.get("architecture", "vit_tiny_patch16_384"))
+        num_classes = int(config_data.get("num_classes", 2))
+        label_names = [str(x).lower() for x in config_data.get("label_names", ["nsfw", "sfw"])]
+        pretrained_cfg = config_data.get("pretrained_cfg", {})
+
+        model = timm.create_model(architecture, pretrained=False, num_classes=num_classes).eval()
+        state_dict = load_file(str(NSFW_WEIGHTS_PATH), device="cpu")
+        model.load_state_dict(state_dict, strict=False)
+
+        # Use local config for preprocessing so inference does not depend on remote metadata.
+        model.pretrained_cfg = {**getattr(model, "pretrained_cfg", {}), **pretrained_cfg, "label_names": label_names}
+
+        data_config = timm.data.resolve_model_data_config(model)
+        transforms = timm.data.create_transform(**data_config, is_training=False)
+
+        _nsfw_runtime = {
+            "torch": torch,
+            "Image": Image,
+            "model": model,
+            "transforms": transforms,
+            "label_names": [str(x).lower() for x in label_names],
+        }
+        return _nsfw_runtime
+    except Exception as exc:
+        _nsfw_error = str(exc)
+        return None
+
+
+def _classify_sensitive_with_model(file_path: str) -> tuple[bool, float, str] | None:
+    runtime = _load_nsfw_runtime()
+    if runtime is None:
+        return None
+
+    torch = runtime["torch"]
+    Image = runtime["Image"]
+    model = runtime["model"]
+    transforms = runtime["transforms"]
+    label_names = runtime["label_names"]
+
+    with Image.open(file_path) as img:
+        img = img.convert("RGB")
+        with torch.no_grad():
+            output = model(transforms(img).unsqueeze(0)).softmax(dim=-1).cpu()[0]
+
+    scores = [float(x) for x in output.tolist()]
+    nsfw_score = _extract_nsfw_score(scores, label_names)
+    return (nsfw_score >= NSFW_THRESHOLD, nsfw_score, "timm-marqo-nsfw")
+
+
+def _extract_nsfw_score(scores: list[float], labels: list[str]) -> float:
+    for idx, label in enumerate(labels):
+        if "nsfw" in label:
+            return scores[idx]
+    if len(scores) >= 2:
+        return scores[1]
+    return scores[0] if scores else 0.0