src/image_tagger.py

from __future__ import annotations
import csv
from dataclasses import dataclass, field
from typing import Optional, Union

try:
    from PIL import ImageOps
except ImportError:  # Pillow is required by app; guard anyway
    ImageOps = None

# Base image-processing deps are required by the app; heavy ML deps remain
# optional so the Tagger tab can degrade gracefully on lightweight installs.
import numpy as np
from PIL import Image

_TAGGER_DEPS_OK = True
try:
    import timm
    import torch
    from huggingface_hub import hf_hub_download
    from huggingface_hub.utils import HfHubHTTPError
    from timm.data import create_transform, resolve_data_config
    from torch import nn
    from torch.nn import functional as F
except Exception:  # pragma: no cover - depends on environment
    _TAGGER_DEPS_OK = False
    timm = torch = None
    hf_hub_download = HfHubHTTPError = create_transform = resolve_data_config = nn = F = None

_REPO_ID = "SmilingWolf/wd-eva02-large-tagger-v3"
_tagger_instance = None

_ENABLED_PKG = {"torch", "timm", "huggingface_hub"}


def _tagger_enabled() -> bool:
    """Respect WHYX_ENABLE_TAGGER (default on). Set to 0/false to skip the
    multi-GB model download on lightweight deployments."""
    import os

    val = os.environ.get("WHYX_ENABLE_TAGGER", "1").strip().lower()
    return val not in ("0", "false", "no", "off")


def _tags_to_caption(tags: list[str]) -> str:
    return ", ".join(tags)


@dataclass
class _LabelData:
    names: list[str] = field(default_factory=list)
    rating: list[int] = field(default_factory=list)
    general: list[int] = field(default_factory=list)
    character: list[int] = field(default_factory=list)
    copyright: list[int] = field(default_factory=list)


def _load_labels(repo_id: str) -> _LabelData:
    try:
        csv_path = hf_hub_download(repo_id=repo_id, filename="selected_tags.csv")
    except HfHubHTTPError as e:
        raise FileNotFoundError(f"selected_tags.csv failed to download from {repo_id}") from e
    labels = _LabelData()
    with open(csv_path, encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for idx, row in enumerate(reader):
            labels.names.append(row["name"])
            cat = int(row.get("category", "0") or 0)
            if cat == 9:
                labels.rating.append(idx)
            elif cat == 0:
                labels.general.append(idx)
            elif cat == 4:
                labels.character.append(idx)
            elif cat == 3:  # copyright (franchise) — previously discarded
                labels.copyright.append(idx)
    return labels


def _ensure_rgb(image: Image.Image) -> Image.Image:
    if image.mode not in ("RGB", "RGBA"):
        image = image.convert("RGBA") if "transparency" in image.info else image.convert("RGB")
    if image.mode == "RGBA":
        canvas = Image.new("RGBA", image.size, (255, 255, 255))
        canvas.alpha_composite(image)
        image = canvas.convert("RGB")
    return image


def _to_pil_image(image) -> Image.Image:
    if isinstance(image, Image.Image):
        img = _ensure_rgb(image)
        return ImageOps.exif_transpose(img) if ImageOps else img
    if np is None:
        raise RuntimeError("numpy is required for image tagging")
    arr = np.asarray(image)
    # Accept bytes / file-like objects (e.g. from API calls or older Gradio versions)
    if arr.ndim == 0 or arr.dtype == object:
        # Accept bytes / bytearray / memoryview / any file-like with .read()
        buf = image if isinstance(image, (bytes, bytearray, memoryview)) else None
        if buf is None and hasattr(image, "read"):
            buf = image.read()
        if buf is None:
            raise ValueError("Unsupported image payload type")
        from io import BytesIO
        arr = np.array(Image.open(BytesIO(bytes(buf))).convert("RGB"))
    if arr.ndim == 2:
        arr = np.stack([arr] * 3, axis=-1)
    elif arr.ndim == 3 and arr.shape[2] == 1:
        arr = np.repeat(arr, 3, axis=2)
    elif arr.ndim == 3 and arr.shape[2] == 4:
        arr = arr[:, :, :4]
    elif arr.ndim != 3 or arr.shape[2] not in (3, 4):
        raise ValueError(f"Unsupported image shape for tagger: {arr.shape}")

    if np.issubdtype(arr.dtype, np.floating):
        scale = 255.0 if arr.max(initial=0) <= 1.0 else 1.0
        arr = np.clip(arr * scale, 0, 255).astype("uint8")
    else:
        arr = np.clip(arr, 0, 255).astype("uint8")

    mode = "RGBA" if arr.shape[2] == 4 else "RGB"
    img = _ensure_rgb(Image.fromarray(arr, mode=mode))
    return ImageOps.exif_transpose(img) if ImageOps else img


def _pad_square(image: Image.Image) -> Image.Image:
    px = max(image.size)
    canvas = Image.new("RGB", (px, px), (255, 255, 255))
    canvas.paste(image, ((px - image.width) // 2, (px - image.height) // 2))
    return canvas


class ImageTagger:
    def __init__(self, repo_id: str = _REPO_ID):
        self._repo_id = repo_id
        self._model: Optional[nn.Module] = None
        self._labels: Optional[_LabelData] = None
        self._transform = None
        self._device: Optional[torch.device] = None
        self._loaded = False

    def ensure_loaded(self):
        if self._loaded:
            return
        if not _TAGGER_DEPS_OK:
            raise RuntimeError("Image tagger dependencies (torch/timm) are not installed.")
        if not _tagger_enabled():
            raise RuntimeError("Image tagger is disabled (WHYX_ENABLE_TAGGER=0).")
        self._model = timm.create_model("hf-hub:" + self._repo_id).eval()
        state_dict = timm.models.load_state_dict_from_hf(self._repo_id)
        self._model.load_state_dict(state_dict)
        self._labels = _load_labels(self._repo_id)
        self._transform = create_transform(**resolve_data_config(self._model.pretrained_cfg, model=self._model))
        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self._model = self._model.to(self._device).cpu()  # always on CPU for HF Spaces free tier
        self._loaded = True

    @property
    def loaded(self) -> bool:
        return self._loaded

    @property
    def available(self) -> bool:
        return _TAGGER_DEPS_OK and _tagger_enabled()

    def tag_image(
        self,
        image: "np.ndarray",
        gen_threshold: float = 0.35,
        char_threshold: float = 0.75,
    ) -> dict:
        self.ensure_loaded()
        pil_img = _to_pil_image(image)
        pil_img = _pad_square(pil_img)
        inputs = self._transform(pil_img).unsqueeze(0)
        inputs = inputs[:, [2, 1, 0]]  # BGR

        device = self._device or torch.device("cpu")
        with torch.inference_mode():
            if device.type != "cpu":
                inputs = inputs.to(device)
            outputs = self._model.forward(inputs)
            outputs = F.sigmoid(outputs)

        probs = outputs.squeeze(0)
        return self._probs_to_dict(probs, gen_threshold, char_threshold)

    def _probs_to_dict(self, probs, gen_threshold: float, char_threshold: float) -> dict:
        named = dict(zip(self._labels.names, probs.tolist()))

        ratings = {self._labels.names[i]: named[self._labels.names[i]] for i in self._labels.rating}
        ratings = {k: round(v, 4) for k, v in ratings.items()}

        def _above(idxs, thresh):
            out = {}
            for i in idxs:
                name = self._labels.names[i]
                s = named[name]
                if s >= thresh:
                    out[name] = round(s, 4)
            return dict(sorted(out.items(), key=lambda x: -x[1]))

        gen_tags = _above(self._labels.general, gen_threshold)
        char_tags = _above(self._labels.character, char_threshold)
        copyright_tags = _above(self._labels.copyright, 0.50)

        caption_names = list(gen_tags.keys()) + list(char_tags.keys())
        caption = ", ".join(caption_names)
        taglist = caption.replace("_", " ").replace("(", "(").replace(")", ")")

        return {
            "caption": caption,
            "taglist": taglist,
            "ratings": ratings,
            "characters": char_tags,
            "copyright": copyright_tags,
            "general": gen_tags,
        }


def get_tagger() -> ImageTagger:
    global _tagger_instance
    if _tagger_instance is None:
        _tagger_instance = ImageTagger()
    return _tagger_instance