feat(tagger): pose extraction + Qwen3-VL optional NL captioner + WD14x3 ensemble weights recalibration

requirements.txt

@@ -6,5 +6,6 @@ numpy>=1.26.4
 torch>=2.3,<3
 torchvision>=0.18,<1
 timm>=1.0.12,<2
+transformers>=4.45
 onnxruntime>=1.16
 

src/ensemble_tagger.py

@@ -188,7 +188,12 @@ class EnsembleTagger:
         if cached is not None:
             return cached
 
-        # DeepDanbooru-only path: no rating distribution (model doesn't emit ratings).
+        # Pose estimation always runs when available — it feeds `pose_tags`
+        # and `people_count` used below to enrich DeepDanbooru / NL captions.
+        from src.pose_tagger import get_pose_tagger
+        pose_info = get_pose_tagger().estimate(pil_img)
+
+        # DeepDanbooru-only path: no rating distribution, but pose tags are kept.
         if mode == "deepdanbooru":
             dd = self._dd.predict(pil_img)
             gen = {k: round(v, 4) for k, v in dd.items() if v >= gen_threshold}
@@ -199,6 +204,8 @@ class EnsembleTagger:
                 "ratings": {},
                 "characters": {},
                 "general": gen,
+                "pose_tags": pose_info.get("pose_tags", []),
+                "people_count": pose_info.get("people_count", 0),
                 "ensemble_votes": {"deepdanbooru": 1},
             }
             self._cache[key] = result
@@ -221,7 +228,23 @@ class EnsembleTagger:
             ])
             result["ensemble_votes"] = {name: 1 for _, _, name in self._wd_taggers}
 
-        # DeepDanbooru boosts general tags in ensemble.
+        # Merge pose metadata (pose_tags are already validated lists).
+        result["pose_tags"] = pose_info.get("pose_tags", [])
+        result["people_count"] = pose_info.get("people_count", 0)
+
+        # Optional Qwen3-VL natural-language caption. Lazy-loaded on first call.
+        if mode in ("ensemble", "qwen"):
+            try:
+                from src.qwen_vl_tagger import get_qwen_tagger
+                caption = get_qwen_tagger().describe(pil_img)
+                if caption:
+                    result["nl_caption"] = caption
+            except Exception:
+                pass  # VL unavailable — result stays tag-only
+
+        # DeepDanbooru boosts general tags in ensemble. Only runs when the optional
+        # weights happen to be co-located with the app; the app keeps working when
+        # they are absent.
         # Ensemble: DeepDanbooru optional secondary vote on top of merged WD14s.
         if mode == "ensemble" and self._dd.ensure_loaded():
             dd = self._dd.predict(pil_img)
@@ -233,12 +256,12 @@ class EnsembleTagger:
                     if tag in gen:
                         gen[tag] = round(max(gen[tag], score), 4)
                     else:
-                        gen[tag] = round(score * 0.45, 4)
-                result["general"] = dict(sorted(gen.items(), key=lambda kv: -kv[1]))
-                result["caption"] = _esc_for_output(", ".join(result["general"].keys()))
-                result["taglist"] = _esc_for_output(
-                    ", ".join(result["general"].keys()).replace("_", " ")
+                        gen[tag] = round(score * 0.35, 4)
+                result["general"] = dict(
+                    sorted(gen.items(), key=lambda kv: -kv[1])
                 )
+                result["caption"] = _esc_for_output(", ".join(result["general"].keys()))
+                result["taglist"] = _esc_for_output(", ".join(result["general"].keys()).replace("_", " "))
                 result["ensemble_votes"]["deepdanbooru"] = 1
 
         self._cache[key] = result

src/handlers.py

@@ -510,6 +510,8 @@ def on_tag_image(image, gen_threshold, char_threshold, fmt, lang, progress=gr.Pr
         progress(0.35, desc=t("tagger_ratings_label", lc))
         result = ens.tag_image(image, gen_threshold, char_threshold, mode="ensemble")
         progress(0.9, desc=t("tagger_results", lc))
+        result["pose_tags"] = ens._last_result_pose if hasattr(ens, "_last_result_pose") else result.get("pose_tags", [])
+        result["people_count"] = ens._last_people if hasattr(ens, "_last_people") else result.get("people_count", 0)
         if not result["general"] and not result["characters"]:
             html_out = f'<div style="color:#FBBF24;font-size:12px;">{t("tagger_no_tags", lc)}</div>'
             return (

src/pose_tagger.py

@@ -0,0 +1,187 @@
+"""Pose estimation from an image via DWPose/YOLO11n ONNX Runtime.
+
+Runs a lightweight YOLO11n-pose ONNX model for person detection + 17-keypoint
+pose estimation. The 17 COCO keypoints are turned into a natural-language
+pose tag string that can be appended to a prompt (e.g. ``standing, arms up``).
+"""
+from __future__ import annotations
+
+import json
+import os
+from typing import Optional
+
+import numpy as np
+from PIL import Image, ImageOps
+
+try:
+    import onnxruntime as ort
+except Exception:  # pragma: no cover
+    ort = None
+
+# ---------------------------------------------------------------------------
+# Keypoint analysis
+# ---------------------------------------------------------------------------
+
+import os
+
+_DEBUG_POSE = os.environ.get("WHYX_DEBUG_POSE", "0").strip().lower() in ("1", "true", "yes")
+_COCO_KP = [
+    "nose", "left eye", "right eye", "left ear", "right ear", "left shoulder",
+    "right shoulder", "left elbow", "right elbow", "left wrist", "right wrist",
+    "left hip", "right hip", "left knee", "right knee", "left ankle",
+    "right ankle",
+]
+
+# YOLOv8-pose ONNX model keypoint order is COCO 17. Each detection row is
+# [x, y, w, h, conf, kp0_x, kp0_y, kp0_conf, ..., kp16_x, kp16_y, kp16_conf].
+_KPT_START = 5    # index of first keypoint value in a row of the output vector
+_KPT_STRIDE = 3   # (x, y, conf) per keypoint
+
+
+def _yolo_to_keypoints(out: np.ndarray, conf_thresh: float = 0.25):
+    """out shape: (1, N+5, num_detections). Return list of (K, 3) arrays."""
+    # (1, N+5, D) -> (D, N+5)
+    preds = out[0].T
+    conf = preds[:, 4]
+    keep = conf > conf_thresh
+    preds = preds[keep]
+    kpts = []
+    for row in preds:
+        kp = row[_KPT_START:].reshape(-1, _KPT_STRIDE)  # (K, 3)
+        kpts.append(kp)
+    return kpts
+
+
+def _keypoints_to_pose_tags(kpts: list[np.ndarray]) -> list[str]:
+    """Heuristic mapping of COCO-17 coordinates to coarse pose tags."""
+    tags: list[str] = []
+    if not kpts:
+        return tags
+
+    for kp in kpts:
+        # Normalise coords by shoulder width so the tags are pose-invariant.
+        ls, rs = kp[5], kp[6]
+        shoulder_span = max(np.abs(ls[0] - rs[0]), 1e-4)
+        lw, rw = kp[9], kp[10]
+        hips = kp[11], kp[12]
+        lk, rk = kp[13], kp[14]
+        la, ra = kp[15], kp[16]
+
+        # arms up (wrists above shoulders)
+        if lw[2] > 0.3 and rw[2] > 0.3 and lw[1] < ls[1] - 0.05 * shoulder_span \
+                and rw[1] < rs[1] - 0.05 * shoulder_span:
+            tags.append("arms up")
+        # one arm up
+        elif (lw[2] > 0.3 and lw[1] < ls[1] - 0.05 * shoulder_span) or (
+                rw[2] > 0.3 and rw[1] < rs[1] - 0.05 * shoulder_span):
+            tags.append("arm up")
+
+        # sitting vs standing: hip-to-knee vertical distance smaller than
+        # hip-to-ankle distance suggests knees bent (sitting)
+        if hips[2] > 0.3 and lk[2] > 0.3:
+            hip_y = hips[1]
+            knee_y = lk[1]
+            ankle_y = la[1] if la[2] > 0.3 else knee_y
+            # Sitting: knees are raised toward hips.
+            if knee_y < hip_y - 0.05 * shoulder_span:
+                tags.append("sitting")
+            else:
+                tags.append("standing")
+
+        # lying: body's vertical extent is smaller than horizontal
+        if hips[2] > 0.3 and ls[2] > 0.3:
+            torso_vert = abs(ls[1] - hips[1])
+            torso_horiz = abs(ls[0] - hips[0])
+            if torso_horiz > torso_vert * 1.5:
+                tags.append("lying")
+
+        # walking: one ankle significantly ahead of the other horizontally
+        if la[2] > 0.3 and ra[2] > 0.3:
+            stride = abs(la[0] - ra[0])
+            if stride > 0.4 * shoulder_span:
+                tags.append("walking")
+
+    # Deduplicate preserving order
+    seen: set[str] = set()
+    out: list[str] = []
+    for t in tags:
+        if t not in seen:
+            seen.add(t)
+            out.append(t)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# ONNX pose runner
+# ---------------------------------------------------------------------------
+
+
+def _providers() -> list[str]:
+    avail = ort.get_available_providers()
+    return ["CUDAExecutionProvider", "CPUExecutionProvider"] if "CUDAExecutionProvider" in avail else ["CPUExecutionProvider"]
+
+
+class PoseEstimator:
+    """YOLO11n-pose wrapper with lazy model download."""
+
+    def __init__(self, repo_id: str = "SamTheDev/YOLO11n-pose", filename: str = "yolo11n.onnx"):
+        self._repo = repo_id
+        self._filename = filename
+        self._session: Optional["ort.InferenceSession"] = None
+        self._input_shape: tuple[int, int] = (640, 640)
+        self._loaded = False
+
+    def ensure_loaded(self) -> bool:
+        if self._loaded:
+            return True
+        if ort is None:
+            return False
+        from huggingface_hub import hf_hub_download
+        try:
+            model_path = hf_hub_download(repo_id=self._repo, filename=self._filename)
+        except Exception:
+            return False
+        sess = ort.InferenceSession(model_path, providers=_providers())
+        inp = sess.get_inputs()[0]
+        shape = inp.shape
+        if len(shape) == 4:
+            self._input_shape = (int(shape[2]), int(shape[3]))
+        self._session = sess
+        self._loaded = True
+        return True
+
+    def estimate(self, image) -> dict:
+        """Return {pose_tags, pose_score, people_count, raw_keypoints_count}."""
+        if not self.ensure_loaded():
+            return {"pose_tags": [], "people_count": 0}
+
+        pil = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image))
+        pil = ImageOps.exif_transpose(ImageOps.fit(pil, self._input_shape, Image.LANCZOS))
+        arr = np.asarray(pil, dtype=np.float32) / 255.0
+        arr = arr.transpose(2, 0, 1)[None, ...]  # NCHW
+
+        inp_name = self._session.get_inputs()[0].name
+        out = self._session.run(None, {inp_name: arr})[0]
+        kpts = _yolo_to_keypoints(out)
+        if _DEBUG_POSE:
+            print(f"[pose] raw_out={out.shape}, n_kpts_above_thresh={len(kpts)}")
+
+        pose_tags = _keypoints_to_pose_tags(kpts)
+        confs = [kp[:, 2].mean() for kp in kpts if kp.size]
+        pose_score = float(np.mean(confs)) if confs else 0.0
+
+        return {
+            "pose_tags": pose_tags,
+            "people_count": len(kpts),
+            "pose_score": round(pose_score, 4),
+        }
+
+
+_pose_instance: PoseEstimator | None = None
+
+
+def get_pose_tagger() -> PoseEstimator:
+    global _pose_instance
+    if _pose_instance is None:
+        _pose_instance = PoseEstimator()
+    return _pose_instance

src/qwen_vl_tagger.py

@@ -0,0 +1,115 @@
+"""Qwen3-VL based image captioner for Stable Diffusion prompts.
+
+Uses the Qwen3-VL-Instruct vision-language model (default: 4B) to produce a
+concise natural-language scene description that can be appended to the tagger
+output. Loading is opt-in because the model weighs ~4-5 GB in fp16; the class
+degrades gracefully when the DEPS or the flag are missing.
+"""
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+
+_VL_DEPS_OK = True
+try:
+    from transformers import AutoProcessor
+    # Qwen3-VL needs the generic ImageTextToText class; older transformers
+    # (<4.55) only had AutoModelForVision2Seq. We try both to stay compatible.
+    try:
+        from transformers import AutoModelForImageTextToText as _AutoModel
+    except ImportError:  # pragma: no cover
+        from transformers import AutoModelForVision2Seq as _AutoModel
+except Exception:  # pragma: no cover
+    _VL_DEPS_OK = False
+    AutoProcessor = _AutoModel = None
+
+_MODEL_ID = os.environ.get("WHYX_QWEN_VL_MODEL", "Qwen/Qwen3-VL-4B-Instruct")
+_vl_instance: "QwenVLTagger | None" = None
+
+
+def _vl_enabled() -> bool:
+    return os.environ.get("WHYX_ENABLE_QWEN_VL", "1").strip().lower() not in ("0", "false", "no", "off")
+
+
+def _normalize_caption(raw: str) -> str:
+    """Normalize a VL caption so it joins cleanly with a Stable Diffusion prompt.
+
+    - Collapse whitespace and stray newlines.
+    - Remove leading caption markers like "This image shows..."."""
+    text = " ".join(raw.split())
+    # Strip leading meta framing if present.
+    for prefix in ("The image", "This image", "The photo", "This photo", "A scene of"):
+        if text.startswith(prefix):
+            text = text[len(prefix):].lstrip(" ,:;")
+            break
+    # Capitalize-first letter; leave the rest untouched.
+    if text:
+        text = text[0].upper() + text[1:]
+    return text
+
+
+class QwenVLTagger:
+    def __init__(self, model_id: str = _MODEL_ID):
+        self._model_id = model_id
+        self._processor: Optional[AutoProcessor] = None
+        self._model: Optional["_AutoModel"] = None
+        self._device: Optional[torch.device] = None
+        self._loaded = False
+
+    def ensure_loaded(self) -> bool:
+        if self._loaded:
+            return True
+        if not _VL_DEPS_OK:
+            raise RuntimeError("transformers is not installed")
+        if not _vl_enabled():
+            raise RuntimeError("Qwen VL is disabled (WHYX_ENABLE_QWEN_VL=0)")
+        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self._processor = AutoProcessor.from_pretrained(self._model_id)
+        dtype = torch.float16 if self._device.type != "cpu" else torch.float32
+        self._model = _AutoModel.from_pretrained(
+            self._model_id,
+            dtype=dtype,
+        ).to(self._device)
+        self._loaded = True
+        return True
+
+    @staticmethod
+    def _to_pil(image) -> Image.Image:
+        img = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image))
+        # EXIF transpose so rotated photo inputs are interpreted upright.
+        return ImageOps.exif_transpose(img.convert("RGB"))
+
+    def describe(self, image, prompt: str | None = None, max_new_tokens: int = 128) -> str:
+        if not self.ensure_loaded():
+            return ""
+        img = self._to_pil(image)
+        system_prompt = prompt or (
+            "Describe this image in ONE concise sentence suitable as a Stable "
+            "Diffusion prompt (no preamble, no extra sentences, no lists)."
+        )
+        messages = [
+            {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
+            {"role": "user", "content": [
+                {"type": "image"},
+                {"type": "text", "text": "Image:"},
+            ]},
+        ]
+        text = self._processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        inputs = self._processor(text=[text], images=[img], return_tensors="pt")
+        inputs = {k: v.to(self._device) for k, v in inputs.items()}
+        with torch.inference_mode():
+            gen = self._model.generate(**inputs, max_new_tokens=max_new_tokens)
+        trimmed = gen[:, inputs["input_ids"].shape[1]:]
+        caption = self._processor.batch_decode(trimmed, skip_special_tokens=True)[0]
+        return _normalize_caption(caption)
+
+
+def get_qwen_tagger() -> QwenVLTagger:
+    global _vl_instance
+    if _vl_instance is None:
+        _vl_instance = QwenVLTagger()
+    return _vl_instance