From d63d680ae8b8b82921ee6aaafbd6e0141a926c9d Mon Sep 17 00:00:00 2001 From: christian1298 Date: Sat, 17 Oct 2026 23:52:41 +0200 Subject: [PATCH] Add TagAI plugin Tag images and videos with locally hosted AI models (WD taggers, JoyTag, YOLO NSFW classifiers). Adds 6 Stash tasks for tagging images and scenes, with optional scene-marker creation for video tagging. Runtime requirements: see plugins/tagAI/README.md (pip deps installed into Stash's Python, ffmpeg for video tagging, ONNX model setup per Models/README.md). --- plugins/tagAI/.gitignore | 12 + plugins/tagAI/Models/README.md | 254 +++ plugins/tagAI/README.md | 569 ++++++ plugins/tagAI/_api.key | 1 + plugins/tagAI/requirements.txt | 13 + plugins/tagAI/tagAI.py | 33 + plugins/tagAI/tagAI.yml | 125 ++ plugins/tagAI/tagai_core/__init__.py | 0 plugins/tagAI/tagai_core/cli.py | 229 +++ plugins/tagAI/tagai_core/config.py | 286 +++ plugins/tagAI/tagai_core/debug.py | 130 ++ plugins/tagAI/tagai_core/frames.py | 124 ++ plugins/tagAI/tagai_core/images.py | 87 + plugins/tagAI/tagai_core/inference.py | 403 ++++ plugins/tagAI/tagai_core/jobs.py | 102 + plugins/tagAI/tagai_core/markers.py | 263 +++ plugins/tagAI/tagai_core/models.py | 423 +++++ plugins/tagAI/tagai_core/postprocess.py | 348 ++++ plugins/tagAI/tagai_core/preprocess.py | 125 ++ plugins/tagAI/tagai_core/scenes.py | 168 ++ plugins/tagAI/tagai_core/stash_client.py | 157 ++ plugins/tagAI/tagai_core/stashlog.py | 42 + plugins/tagAI/tagai_core/tags.py | 61 + plugins/tagAI/tests/test_tagai.py | 2205 ++++++++++++++++++++++ 24 files changed, 6160 insertions(+) create mode 100755 plugins/tagAI/.gitignore create mode 100755 plugins/tagAI/Models/README.md create mode 100755 plugins/tagAI/README.md create mode 100755 plugins/tagAI/_api.key create mode 100755 plugins/tagAI/requirements.txt create mode 100755 plugins/tagAI/tagAI.py create mode 100755 plugins/tagAI/tagAI.yml create mode 100755 plugins/tagAI/tagai_core/__init__.py create mode 100755 plugins/tagAI/tagai_core/cli.py create mode 100755 plugins/tagAI/tagai_core/config.py create mode 100755 plugins/tagAI/tagai_core/debug.py create mode 100755 plugins/tagAI/tagai_core/frames.py create mode 100755 plugins/tagAI/tagai_core/images.py create mode 100755 plugins/tagAI/tagai_core/inference.py create mode 100755 plugins/tagAI/tagai_core/jobs.py create mode 100755 plugins/tagAI/tagai_core/markers.py create mode 100755 plugins/tagAI/tagai_core/models.py create mode 100755 plugins/tagAI/tagai_core/postprocess.py create mode 100755 plugins/tagAI/tagai_core/preprocess.py create mode 100755 plugins/tagAI/tagai_core/scenes.py create mode 100755 plugins/tagAI/tagai_core/stash_client.py create mode 100755 plugins/tagAI/tagai_core/stashlog.py create mode 100755 plugins/tagAI/tagai_core/tags.py create mode 100755 plugins/tagAI/tests/test_tagai.py diff --git a/plugins/tagAI/.gitignore b/plugins/tagAI/.gitignore new file mode 100755 index 00000000..c8d2292d --- /dev/null +++ b/plugins/tagAI/.gitignore @@ -0,0 +1,12 @@ +Models/*.csv +Models/*.onnx +Models/*.part +Models/*-thresholds.json +api.key +tagai-debug.json +tagai-debug-*.json +hf_token +hf.token +huggingface.token +__pycache__/ +*.pyc diff --git a/plugins/tagAI/Models/README.md b/plugins/tagAI/Models/README.md new file mode 100755 index 00000000..2c38b336 --- /dev/null +++ b/plugins/tagAI/Models/README.md @@ -0,0 +1,254 @@ +# Models + +Place the AI model files in this folder. Each model needs the exact file name +listed below. + +Lazy loading means you only need the model(s) for the target(s) you use: + +* Image tasks need one or more **WD taggers** (`.onnx` + `.csv`) or the + **JoyTag** tagger (`.onnx` + `top_tags.txt`). +* Video tasks need the **YOLO classifier** (`.onnx`). + +The `.onnx` and `.csv` files are git-ignored (they are large), so they are not +committed to this repository -- download them yourself. + +## Automatic download + +If a required file is missing when a task runs, TagAI downloads it from +Hugging Face automatically (this can be disabled with `auto_download=false`). +The mapping from model name to repository and file is built in, so normally +you do not have to do anything. + +* The WD taggers and JoyTag are public and download without credentials. +* The `yolo11-nsfw` detector and the `sex-position` classifier are **not** + downloadable (their upstream repositories ship only `.pt`/`.bin` + checkpoints). Export their ONNX once, as described below. +* The YOLO classifier is **gated** (see below) and needs a token. Provide one + by setting the `HF_TOKEN` environment variable (`HUGGING_FACE_HUB_TOKEN` and + `HUGGINGFACE_TOKEN` also work), by passing the `hf_token` task argument, or + by placing the token in a file named `hf_token` (or `hf.token` / + `huggingface.token`) next to `tagAI.py`. +* For the YOLO classifier TagAI also tries to fetch the model's published + per-class threshold file (best effort). If it cannot be downloaded, the + built-in defaults are used instead. + +If a download fails you can always fetch the files manually and drop them here +with the exact names below; the automatic download only fills in what is +missing. + +--- + +## WD taggers (images) + +Multi-label taggers from [SmilingWolf](https://huggingface.co/SmilingWolf). +For each model download `model.onnx` and `selected_tags.csv`, then rename them +so the base name matches the model name. + +### wd-v1-4-moat-tagger-v2 (default for images) + +* https://huggingface.co/SmilingWolf/wd-v1-4-moat-tagger-v2/resolve/main/model.onnx?download=true +* https://huggingface.co/SmilingWolf/wd-v1-4-moat-tagger-v2/resolve/main/selected_tags.csv?download=true + +Rename to: + +* `wd-v1-4-moat-tagger-v2.onnx` +* `wd-v1-4-moat-tagger-v2.csv` + +### wd-convnext-tagger-v3 + +* https://huggingface.co/SmilingWolf/wd-convnext-tagger-v3/resolve/main/model.onnx?download=true +* https://huggingface.co/SmilingWolf/wd-convnext-tagger-v3/resolve/main/selected_tags.csv?download=true + +Rename to: + +* `wd-convnext-tagger-v3.onnx` +* `wd-convnext-tagger-v3.csv` + +### wd-v1-4-vit-tagger + +* https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger/resolve/main/model.onnx?download=true +* https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger/resolve/main/selected_tags.csv?download=true + +Rename to: + +* `wd-v1-4-vit-tagger.onnx` +* `wd-v1-4-vit-tagger.csv` + +### wd-swinv2-tagger-v3 + +An updated tagger that recognises **10,861** tags (the v1-4 series has 9,083). + +* https://huggingface.co/SmilingWolf/wd-swinv2-tagger-v3/resolve/main/model.onnx?download=true +* https://huggingface.co/SmilingWolf/wd-swinv2-tagger-v3/resolve/main/selected_tags.csv?download=true + +Rename to: + +* `wd-swinv2-tagger-v3.onnx` +* `wd-swinv2-tagger-v3.csv` + +### wd-eva02-large-tagger-v3 + +The largest and most accurate WD v3 tagger (10,861 tags). It is noticeably +slower and uses more memory than the other taggers. + +* https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3/resolve/main/model.onnx?download=true +* https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3/resolve/main/selected_tags.csv?download=true + +Rename to: + +* `wd-eva02-large-tagger-v3.onnx` +* `wd-eva02-large-tagger-v3.csv` + +--- + +## JoyTag (images, photographic) + +[JoyTag](https://huggingface.co/fancyfeast/joytag) is a multi-label tagger +(ViT-B/16, 448x448, 5,000+ Danbooru tags) that was additionally trained on +manually tagged **photographic** images, so it tends to work better on real +photos than the WD taggers. + +* https://huggingface.co/fancyfeast/joytag/resolve/main/model.onnx?download=true +* https://huggingface.co/fancyfeast/joytag/resolve/main/top_tags.txt?download=true + +Rename to: + +* `joytag.onnx` +* `joytag-top_tags.txt` + +`top_tags.txt` holds one tag per line, in the model's output order. JoyTag +outputs raw logits; TagAI applies a sigmoid and keeps every tag above the +`joytag_threshold` argument (default `0.4`). Images are padded to a white +square, resized to 448x448 and CLIP-normalised before inference. + +--- + +## YOLO classifier (videos) + +The video pipeline uses +[`likewendy/yolo26n-cls-porn`](https://huggingface.co/likewendy/yolo26n-cls-porn), +a single-label classifier over 35 adult classes. + +> **This model is gated.** You must be logged in to Hugging Face and have +> accepted the model's terms before you can download it. + +1. Log in to Hugging Face and accept the model's license/terms on its page. +2. Download the weights (`weights/best.onnx` from the repository). +3. Place the file here and rename it to: + + * `yolo26n-cls-porn.onnx` + +> **Note:** the exported model already contains a Softmax layer, so its output +> is a probability distribution. TagAI detects and handles both softmaxed +> output and raw logits, so either export works. Do not rename the classes. + +### Per-class thresholds + +The model ships a published per-class confidence threshold file +(`class35_conf_thresholds_by_name.json`); TagAI downloads it here as +`yolo26n-cls-porn-thresholds.json` when possible and otherwise uses the same +values bundled into the code. + +To override them, place a JSON object mapping class name to a threshold here: + +* `yolo26n-cls-porn-thresholds.json` + +```json +{ + "Kissing": 0.6, + "Pants": 0.85, + "default": 0.5 +} +``` + +When present, a frame's top class is accepted only if it clears its own +threshold (falling back to `default`, then to the `yolo_threshold` argument). +Keys may be given in Chinese or English (the model's own names are Chinese and +are translated automatically). + +--- + +## JoyTag (images) + +See the JoyTag section above; it lives in this folder as `joytag.onnx` and +`joytag-top_tags.txt`. + +--- + +## YOLO11-NSFW detector (videos) + +`yolo11-nsfw` is an object detector that finds explicit acts and body parts +(10 classes: person, breast, vulva, butt, male, penis, anal, vaginal, blowjob, +handjob). Unlike the classifiers it can report several classes for one frame. + +TagAI **cannot download this model for you**: the upstream repository +([`Felldude/Yolo11_NSFW_Small`](https://huggingface.co/Felldude/Yolo11_NSFW_Small)) +only publishes a `.pt` checkpoint and its licence is **AGPL-3.0**, so the ONNX +has to be exported once by hand: + +```sh +# 1. download the checkpoint +curl -L -o yolo_nsfw_s.pt \ + https://huggingface.co/Felldude/Yolo11_NSFW_Small/resolve/main/yolo_nsfw_s.pt + +# 2. export to ONNX (needs `pip install ultralytics`) +yolo export model=yolo_nsfw_s.pt format=onnx imgsz=640 + +# 3. put the result here under the name TagAI expects +mv yolo_nsfw_s.onnx /path/to/tagAI/Models/yolo11-nsfw.onnx +``` + +The exported graph ends at the raw detection head (`[1, 14, 8400]`); TagAI +decodes it itself, filtering boxes with `det_conf` (default `0.25`) and +applying non-maximum suppression with `det_iou` (default `0.45`). Images are +letterboxed to 640x640 with grey padding, exactly like Ultralytics does. + +TagAI ships no weights for this model. If `yolo11-nsfw.onnx` is missing the +task fails with an error pointing here. + +--- + +## sex-position (images and videos) + +`porntech/sex-position` is a single-label resnet50 over eight positions: +`blowjob`, `hardcore`, `titjob`, `handjob`, `pussy-licking`, `fingering`, +`other` and `solo`. The repository ships only `pytorch_model.bin`, so TagAI +cannot download it -- export the ONNX once: + +```bash +python -m venv /tmp/sexpos && . /tmp/sexpos/bin/activate +pip install torch timm onnx +python - <<'PY' +import torch, timm +m = timm.create_model("hf_hub:porntech/sex-position", pretrained=True) +m.eval() +torch.onnx.export( + m, torch.randn(2, 3, 224, 224), "sex-position.onnx", + input_names=["input"], output_names=["output"], + dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}, + opset_version=17, dynamo=False, +) +PY +mv sex-position.onnx /path/to/tagAI/Models/sex-position.onnx +``` + +TagAI expects input `input` with shape `[batch, 3, 224, 224]` and output +`output` with shape `[batch, 8]` **raw logits** (it applies the softmax +itself). Images are resized so the shorter side is 256 (bicubic), +centre-cropped to 224x224 and normalised with the ImageNet mean/std, which +matches the model's own `timm` preprocessing. The weights are MIT licensed +and the model is undefined on non-explicit input. + +## Class names + +The 35 YOLO classes, **in model output order**, are: + +AI Generated, Titfuck, Voyeurism, Public Exposure, Creampie, Anime, Curly Hair, +Double Penetration, Blowjob, Bukkake, Christmas, Maid, Lesbian Strapon, Lesbian +Tribbing, CFNM, Pregnant, Pissing, Handjob, Kissing, Oiled, Squirting, Solo +Male, Male Group, Leather and Latex, Tattoo, Cuckold, Asshole Close-up, Footjob, +Armpit, Rimjob, Pussy Licking, Pants, Interracial, Mirror, Facesitting. + +The model stores these names in Chinese; TagAI reads them from the ONNX +metadata, translates them to English, and falls back to this built-in list if +the metadata is unavailable. diff --git a/plugins/tagAI/README.md b/plugins/tagAI/README.md new file mode 100755 index 00000000..a8e9a0f9 --- /dev/null +++ b/plugins/tagAI/README.md @@ -0,0 +1,569 @@ +# tagAI + +A [Stash](https://stashapp.cc) plugin that tags **Images** and **Videos +(Scenes)** using locally hosted AI models. Nothing is uploaded anywhere -- all +inference runs on your own machine. + +It supports two families of models and lets you choose which to use per task: + +| Target | Model family | Example model | Output | +| -------- | -------------------------------- | -------------------------------------- | ------------------------------------------ | +| `images` | WD / JoyTag multi-label taggers | `wd-v1-4-moat-tagger-v2`, `joytag` | many tags per image | +| `videos` | YOLO classifier / detector + WD | `yolo11-nsfw,sex-position,wd-v1-4-moat-tagger-v2` | one or more tags per video (voted over frames) | + +Models are loaded **lazily**: image jobs never load the YOLO model and video +jobs never load a WD tagger, so only the model you actually need is kept in +memory. + +--- + +## Quick start + +Get it running in five steps. Everything below is explained in more depth +[after the configuration guide](#in-depth-guide). + +### 1. Install the plugin + +Put this whole folder in your Stash community plugins directory, so that the +file `tagAI.py` ends up here: + +``` +/plugins/community/tagAI/tagAI.py +``` + +On a typical Linux install that is `~/.stash/plugins/community/tagAI/`. If Stash +already runs other community plugins, you probably already know where it lives. + +Copy the `tagai_core/` folder alongside `tagAI.py` -- `tagAI.py` is only a thin +launcher and most of the code lives in that package (see +[Project layout](#project-layout) for what is in it). + +### 2. Install the Python dependencies + +Install them into **the same Python that Stash uses to run plugins**: + +```sh +pip install -r requirements.txt +``` + +> ⚠️ This is the single most common setup mistake. "The Python Stash uses" is +> the first `python` on Stash's `PATH`, which is often a dedicated virtualenv -- +> **not** the one your terminal happens to use. Installing into the wrong +> interpreter gives `ModuleNotFoundError: No module named 'PIL'` (or +> `onnxruntime`) when a task runs. + +### 3. Install ffmpeg (only if you want to tag videos) + +`ffmpeg` must be on your `PATH`, or you can point a task at it with the +`ffmpeg_path` argument. Image tagging does not need ffmpeg. + +### 4. Add your model files + +* **Images:** nothing to do. The public WD taggers and JoyTag download + automatically the first time a task needs them. +* **Videos:** the `yolo26n-cls-porn` classifier is gated and needs a Hugging + Face token, and the `yolo11-nsfw` detector / `sex-position` classifier must + be exported to ONNX once by hand. See + [`Models/README.md`](Models/README.md) for the exact steps. + +### 5. Run a task + +In Stash: **Settings → Tasks**, pick one of the `TagAI` tasks (for example +**"Tag All Untagged Images"**) and press **Run**. + +That's it. The first run may download a model and can take a while; later runs +reuse it. + +### Optional: the API key file + +The plugin works out of the box inside Stash because Stash hands it the API +key/session on stdin. If you ever run `tagAI.py` by itself, copy `_api.key` to +`api.key` and put your Stash API key in it. + +--- + +## Configuration (`tagAI.yml`) + +All of the plugin's settings live in **`tagAI.yml`**, in the plugin folder. This +section explains how to read and edit it. A complete list of every available +setting is in the [full argument reference](#full-argument-reference). + +### The shape of the file + +```yaml +name: TagAI # plugin name shown in Stash +version: 1.8 # plugin version +exec: # how Stash launches the plugin + - python + - "{pluginDir}/tagAI.py" +description: Tag images and videos with locally hosted AI models +interface: raw +tasks: # the tasks Stash lists under Settings -> Tasks + - name: "Tag All Untagged Images" + description: "Tag all untagged images (WD tagger)" + defaultArgs: # <-- every setting for this task lives here + target: images + model: wd-v1-4-moat-tagger-v2 + mode: TagAll + flag_tag: "AI_TagMe" + add_tag: "AI_Tagged" + tag_prefix: "AI: " + # ...plus the debug_* and tag_*_list options (see the reference table) +``` + +You normally only ever touch the **`defaultArgs:`** block of a task -- that is +where each task's settings are stored. The rest of the file (name, version, +`exec`, ...) is boilerplate that lets Stash find and run the plugin; leave it +alone. + +### How task settings work + +Each task has its own set of defaults under `defaultArgs:`. When Stash runs a +task it passes those settings to the plugin. **Editing a value in `defaultArgs` +changes the behaviour of that task.** + +To change something, find the task you care about and add or edit the line: + +```yaml + target: images + model: wd-eva02-large-tagger-v3 # changed: use the biggest WD tagger + tag_threshold: "0.6" # changed: require more confidence +``` + +A few rules of thumb: + +* **Values are plain YAML.** Text can be quoted (`"AI: "`) or unquoted + (`images`). Keep the quotes around anything with leading/trailing spaces or a + colon, like the `tag_prefix` value `"AI: "`. +* **Numbers are written as strings** in this file (`"0.5"`, `"10"`) to match the + existing style; the plugin converts them for you. +* **To turn an option off**, set it to an empty value: `markers: ""`. +* **After editing, reload the plugin** in Stash (Settings → Plugins → *Reload*, + or restart Stash) so the new defaults are picked up. + +### The six tasks that ship with the plugin + +| # | Task name | Target | Mode | Model (default) | +| - | ----------------------------- | -------- | ---------- | -------------------------------------------------- | +| 1 | Tag All Untagged Images | `images` | `TagAll` | `wd-v1-4-moat-tagger-v2` | +| 2 | Tag All Images -- Force | `images` | `ForceTag` | `wd-v1-4-moat-tagger-v2` | +| 3 | Tag All Flagged Images | `images` | `Tagged` | `wd-v1-4-moat-tagger-v2` | +| 4 | Tag All Untagged Videos | `videos` | `TagAll` | `yolo11-nsfw,sex-position,wd-v1-4-moat-tagger-v2` | +| 5 | Tag All Videos -- Force | `videos` | `ForceTag` | `yolo11-nsfw,sex-position,wd-v1-4-moat-tagger-v2` | +| 6 | Tag All Flagged Videos | `videos` | `Tagged` | `yolo11-nsfw,sex-position,wd-v1-4-moat-tagger-v2` | + +The `mode` decides which items a task touches: + +* **`TagAll`** -- tag items that do **not** already carry the `add_tag` + (`AI_Tagged`). This is the normal "tag everything new" mode. +* **`ForceTag`** -- tag everything, even items already carrying `add_tag`. + Use it to re-tag after changing thresholds or models. +* **`Tagged`** -- tag only items carrying the `flag_tag` (`AI_TagMe`) -- handy + for tagging a hand-picked set. In Stash, add the `AI_TagMe` tag to the items + you want, then run the "Flagged" task. + +### Common configuration recipes + +**Use a different image model** (in one of the image tasks): + +```yaml + model: wd-eva02-large-tagger-v3 +``` + +**Run several image models at once** and merge their tags: + +```yaml + model: "joytag,wd-eva02-large-tagger-v3" +``` + +**Enable scene markers on a video task** (see +[scene markers](#tagging-exact-moments-scene-markers) below): + +```yaml + markers: ranges # off (default) | ranges | best + marker_max: "200" +``` + +**Tune how confident a model must be:** + +```yaml + tag_threshold: "0.65" # images: WD general tags + det_conf: "0.25" # videos: yolo11-nsfw detections + yolo_min_votes: "2" # videos: frames that must agree +``` + +**Write a debug report to calibrate thresholds:** + +```yaml + debug_log: "tagai-debug.json" +``` + +**Add an entirely new task:** copy an existing task block (from `- name:` down +through its `defaultArgs:`) and change its `name` and settings. No code changes +are needed -- pick a `target` (`images` or `videos`), a `model`, and a `mode`, +and Stash will list it. + +The video tasks that ship with the plugin already run +`yolo11-nsfw,sex-position,wd-v1-4-moat-tagger-v2` at `frame_interval: 5`, +`det_conf: 0.25`, `tag_threshold: 0.65`, `yolo_min_votes: 2` and +`sexpos_margin: 0.15`. The "Tag All Flagged Videos" task additionally has +`markers: ranges` enabled. + +--- + +# In-depth guide + +Everything below is reference material. Read it as you need it. + +## Requirements + +1. **Python dependencies** -- install with `pip install -r requirements.txt`. + Make sure you install them into the **same Python interpreter Stash uses to + run plugins** (the first `python` on Stash's `PATH`, which is not + necessarily the one on your shell). Installing into the wrong interpreter is + the usual cause of `ModuleNotFoundError: No module named 'PIL'` / + `onnxruntime`. +2. **ffmpeg** -- required for video tagging (frame extraction). It must be on + your `PATH`, or provide the `ffmpeg_path` argument on the task. +3. **Models** -- if a model file is missing, TagAI downloads it from Hugging + Face automatically (see [`Models/README.md`](Models/README.md)). The gated + YOLO classifier additionally needs a Hugging Face token once you have + accepted its terms -- set `HF_TOKEN` (or a local `hf_token` file), or pass + the `hf_token` argument. The `yolo11-nsfw` detector and `sex-position` + classifier cannot be downloaded and must be exported to ONNX once by hand. +4. **API key** -- copy `_api.key` to `api.key` and put your Stash API key in it. + (The plugin also accepts the key/session cookie Stash passes on stdin, so + this file is optional when running inside Stash.) + +## Tasks + +Six tasks are registered -- three for images and three for videos. See +[the six tasks](#the-six-tasks-that-ship-with-the-plugin) above for the task +table and the meaning of the `TagAll` / `ForceTag` / `Tagged` modes. + +The video tasks run three models in one pass by default, so every frame +contributes the detector's acts and body parts, the position classifier's tag +and the tagger's Danbooru tags (including named positions such as +`cowgirl position`). + +In all cases the processed item also receives the `add_tag`, and every AI tag is +created automatically if it does not exist yet. AI tags are prefixed with +`tag_prefix` (`AI: ` by default). + +## Choosing a model + +The `model` argument selects which model a task uses. Available models: + +* `wd-v1-4-moat-tagger-v2` (default for images) +* `wd-convnext-tagger-v3` +* `wd-v1-4-vit-tagger` +* `wd-swinv2-tagger-v3` (WD v3, 10,861 tags) +* `wd-eva02-large-tagger-v3` (WD v3, largest and most accurate, slowest) +* `joytag` (photographic multi-label tagger) +* `yolo26n-cls-porn` (35-class single-label classifier, videos) +* `yolo11-nsfw` (YOLO11 NSFW object detector, videos -- see below) +* `sex-position` (porntech sex-position classifier, images and videos) + +The `target` argument chooses the pipeline (`images` or `videos`); it also +determines the default model when `model` is omitted. + +## How tagging works + +### Images (WD taggers) + +The WD models are multi-label: for each image the model returns a probability +per tag. General tags above `tag_threshold` (default `0.5`) are applied. +Character tags (category `4`) are only applied when `tag_characters` is `true` +and above `char_threshold` (default `0.85`). + +The three `wd-*-v3` models (`wd-convnext-tagger-v3`, `wd-swinv2-tagger-v3`, +`wd-eva02-large-tagger-v3`) work exactly the same way but recognise 10,861 tags +instead of 9,083. + +### Images (JoyTag) + +JoyTag is a second multi-label tagger you can select with `model: joytag`. It +was trained on Danbooru plus a manually tagged set of **photographic** images, +so it tends to do better on real photos than on drawn art. Every tag scoring +above `joytag_threshold` (default `0.4`) is applied; there is no separate +character category. + +### Videos (YOLO classifier) + +The YOLO model is single-label *per frame*, but a video usually contains more +than one kind of content, so TagAI votes over many frames and can apply several +tags: + +1. Frames are sampled across the middle 10%-90% of the video's duration with + ffmpeg. By default roughly **one frame every `frame_interval` seconds** + (default `10`), capped at `max_frames` (default `512`). Set + `frame_interval: 0` to instead sample exactly `frame_count` frames (default + `8`) evenly spread over that band. Frames are extracted by `frame_workers` + (default `4`) parallel ffmpeg processes, which hides the seek latency of slow + or networked storage. +2. Every frame is classified. A frame is *accepted* only if the top class clears + its threshold **and** beats the runner-up by at least `yolo_margin` (default + `0.4`). Thresholds are per class: the values published with the model are + bundled in and an override file next to the model is used if present (see + [`Models/README.md`](Models/README.md)); `yolo_threshold` (default `0.5`) is + only the fallback. Frames are classified in batches of `infer_batch_size` + (default `8`), which keeps the GPU busy on long videos. With `yolo_top_k` + above `1` each accepted frame may also report its runner-up classes, but only + when they clear their own threshold. +3. TagAI collects every class that won an accepted frame. With `yolo_multi` + `true` (the default) a class is applied when it wins at least + `yolo_min_votes` accepted frames (default `1`, so a single confident frame is + enough) and, if `yolo_share` is above `0`, at least that fraction of all + accepted frames (default `0` = no share filter). Raise `yolo_min_votes` or + set `yolo_share` to demand more agreement before a tag is applied. If no + class reaches the bar the single strongest class is applied so a video is + never left untagged. Set `yolo_multi: false` to restore the old "one + majority-vote winner" behaviour. + +Class names are stored in the model in Chinese and are automatically translated +to English. + +`yolo11-nsfw` is an **object detector** rather than a single-label classifier: +each frame may contain several boxes, so it can report several acts or body +parts (person, breast, vulva, butt, male, penis, anal, vaginal, blowjob, +handjob) for the same frame. Boxes are filtered with `det_conf` (default +`0.25`) and non-maximum suppressed with `det_iou` (default `0.45`), then the +surviving classes are pooled across frames with the same vote rules as above +(`yolo_multi`, `yolo_min_votes`, `yolo_share`). + +`sex-position` (porntech) is a single-label classifier over eight positions: +`blowjob`, `hardcore` (missionary/doggystyle/cowgirl/spooning), `titjob`, +`handjob`, `pussy-licking`, `fingering`, `other` and `solo`. Its top class is +applied when it reaches `sexpos_threshold` (default `0.5`) and, if +`sexpos_margin` is set, beats the runner-up by that much. On videos the accepted +classes are pooled with the same vote rules as the other video models. It is +undefined on non-explicit input, so it can attach a position tag to a scene that +has none. + +## Running several models at once + +`model` accepts a comma-separated list. Every model is run on the same item in +one pass and the resulting tags are merged -- a tag produced by more than one +model keeps its highest score: + +```yaml + model: "yolo26n-cls-porn,yolo11-nsfw" # videos + model: "yolo11-nsfw,sex-position" # videos + model: "joytag,wd-eva02-large-tagger-v3" # images +``` + +Each model keeps its own options (`det_conf`, `sexpos_threshold`, `yolo_margin`, +...) and every referenced model must be available (downloaded automatically +where possible -- see [`Models/README.md`](Models/README.md)). Every extra model +costs one extra inference pass per item, which is small next to frame extraction +for videos. + +## Tagging exact moments (scene markers) + +Plain Stash tags belong to the whole scene, so a tag search always returns the +full video. If you want to find *the moment* a tag appears, let TagAI create +**scene markers**: Stash's markers carry a start and end time (plus an optional +preview clip), and the Scenes page can filter by marker tag. + +Enable it per task -- it is off by default: + +```yaml + markers: ranges # off (default), ranges or best + marker_pad: "" # seconds added to each end; empty = the sampling interval + marker_max: "200" # marker cap per scene; 0 disables markers + marker_gap: "60" # merge markers of one tag closer than 2*marker_pad + this + marker_whitelist: "" # only these tags may create markers (comma separated) + marker_blacklist: "" # these tags never create markers (comma separated) +``` + +* `ranges` (recommended) -- every sampled frame that accepted a tag becomes a + point, and consecutive points of the same tag are merged into one marker. Two + hits also merge when their padded spans would overlap, so a tag that fires + every other frame does not produce a marker per frame. +* `best` -- a single marker per tag, at its most confident frame. + +Each marker's **primary tag** is the `AI: `-prefixed tag it came from (Stash +requires exactly one primary tag) and the other AI tags that fired on the same +frames are attached as additional marker tags, so you can filter by any of them. + +`marker_gap` merges markers of the *same* tag that are close together, so a tag +which fires on and off for a while becomes one marker per burst instead of one +per sample. Two hits are combined when they are no further apart than +`2 * marker_pad + marker_gap` seconds (`0` = only hits whose padded spans +already touch, which is the old behaviour). With `frame_interval: 5`, +`marker_pad: 2` and `marker_gap: 60`, a tag has to stay quiet for over a minute +before it can produce a second marker. It never merges two different tags. + +`marker_whitelist` and `marker_blacklist` restrict which tags may create +markers, which is how you keep background noise (`blurry`, `1girl`, ...) from +producing marker spam. Both take a comma-separated list of tag names; matching +is case-insensitive, the `AI: ` prefix is optional, and multi-word names such as +`cowgirl position` are kept intact. A non-empty whitelist allows *only* the +listed tags to create markers, and the blacklist always wins over the whitelist. +The filters affect markers only -- the scene still gets all the tags -- and they +also stop a filtered tag from being attached as an extra marker tag. + +`marker_pad` defaults to the median spacing of the sampled frames (for a +`frame_interval` of `5`, that is 5 s), so a marker always spans at least +`2 * marker_pad = 10` s around the frame that fired -- a hit at 60 s becomes a +marker from 55 s to 65 s. That padding is what makes the markers comfortable to +click through, and it is also the merge tolerance in `ranges` mode. + +To search: **Scenes page → Filter → Markers → Marker tags** (add more than one +to require all of them), optionally combined with **Has markers**. Markers also +appear in the scene's *Markers* tab, where clicking one seeks the player to that +moment. + +Notes: + +* Marker previews (the little looping clips under each marker) are rendered by + Stash with ffmpeg; if you generate a lot of markers you may want to turn + marker preview generation down or off in Stash's own settings. +* Re-running a task does not stack markers: a marker is skipped when one with + the same primary tag already exists within `2 * marker_pad` seconds. +* Markers are only created for tags the scene vote accepted (the same thresholds + apply), so tune those first with the debug log below. + +## Filtering the tags + +`tag_whitelist` and `tag_blacklist` decide which tags are written to the item +itself. They work exactly like the marker filters (comma-separated names, +case-insensitive, an optional `AI: ` prefix, the blacklist always wins) but they +act earlier: a filtered tag is never applied to the image or scene, and it can +no longer become a marker either -- so they are a second way to keep marker +noise down. + +```yaml + tag_blacklist: "blurry, 1girl, artist name" # never applied + tag_whitelist: "cowgirl position, doggystyle" # only these are applied +``` + +Both default to empty (= no filtering) and apply to images and videos alike. +With a non-empty whitelist only the listed tags are applied, which is useful +when you want nothing but the acts you care about: everything else the models +produce is dropped before it reaches your library. + +## Calibrating thresholds (debug log) + +Thresholds are guesses until you measure them on your own library. Set +`debug_log` and TagAI writes a JSON report of every tag confidence it saw for +each processed item, so running **one** video with the task you want to tune is +enough: + +```yaml + debug_log: "tagai-debug.json" # a path, or true for tagai-debug.json next to the plugin + debug_floor: "0.05" # ignore scores below this + debug_top: "10" # candidates kept per frame and model +``` + +The report is rewritten after every item (`per_frame` holds one row per sampled +frame): + +```json +{ + "plugin": "TagAI", "format": 1, "generated": "2026-10-17T12:00:00", + "settings": { "det_conf": 0.25, "tag_threshold": 0.5, "...": "..." }, + "items": [ + { "id": 21418, "path": "/media/videos/a.mp4", "duration": 1696.6, "frames": 52, + "models": [ + { "name": "wd-v1-4-moat-tagger-v2", "family": "wd", + "settings": { "tag_threshold": 0.5, "char_threshold": 0.85, "tag_characters": false }, + "per_frame": [ [["sex from behind", 0.814], ["doggystyle", 0.73]], ["..."] ], + "aggregate": [ {"tag": "sex from behind", "score": 0.814, "accepted": true} ], + "tags": ["sex from behind", "doggystyle"] } + ], + "applied_tags": ["AI: sex from behind", "AI: doggystyle"] } + ] +} +``` + +Read the per-model `aggregate` list: every tag the model produced above +`debug_floor`, its best score over the frames and whether the current settings +accepted it. Raise a threshold when a tag with a low score keeps appearing, +lower it when a real tag sits just under the bar. `per_frame` shows whether a +tag fired consistently or on a single lucky frame. The report is written to disk +only -- never to stdout -- and `settings` never contains your `hf_token`. + +## Full argument reference + +Each task declares its defaults in `tagAI.yml`; all of them can be overridden. + +| Argument | Default | Applies to | Meaning | +| ----------------- | ------------------------ | ---------- | ---------------------------------------------------- | +| `target` | `images` | both | `images` or `videos` | +| `model` | per target | both | Model name, or a comma-separated list of names | +| `mode` | `TagAll` | both | `TagAll`, `ForceTag` or `Tagged` | +| `flag_tag` | `AI_TagMe` | both | Tag that marks an item as flagged (mode `Tagged`) | +| `add_tag` | `AI_Tagged` | both | Tag added to every processed item | +| `tag_prefix` | `AI: ` | both | Prefix for all created AI tags | +| `tag_threshold` | `0.5` | images | WD general tag threshold | +| `char_threshold` | `0.85` | images | WD character tag threshold | +| `tag_characters` | `false` | images | Whether to apply WD character tags | +| `joytag_threshold`| `0.4` | images | JoyTag score threshold | +| `frame_interval` | `10` (seconds) | videos | Sample ~one frame every N seconds (0 = use `frame_count`) | +| `max_frames` | `512` | videos | Upper bound on frames sampled per video | +| `frame_count` | `8` | videos | Frames sampled when `frame_interval` is 0 | +| `frame_workers` | `4` | videos | Parallel ffmpeg processes used to extract frames | +| `infer_batch_size`| `8` | videos | Frames per onnxruntime call | +| `yolo_threshold` | `0.5` | videos | Fallback per-frame confidence threshold | +| `yolo_margin` | `0.4` | videos | Minimum top1-top2 confidence gap to accept a frame | +| `yolo_multi` | `true` | videos | Tag every class that wins enough accepted frames | +| `yolo_min_votes` | `1` | videos | Minimum accepted frames for a class to be tagged | +| `yolo_share` | `0` | videos | Minimum fraction of accepted frames (0 = off) | +| `yolo_top_k` | `1` | videos | Classes reported per frame (opt-in recall boost) | +| `det_conf` | `0.25` | videos | Detection confidence threshold (`yolo11-nsfw`) | +| `det_iou` | `0.45` | videos | NMS IoU threshold (`yolo11-nsfw`) | +| `sexpos_threshold`| `0.5` | both | `sex-position` top-1 confidence threshold | +| `sexpos_margin` | `0` (off) | both | Minimum top1-top2 gap for `sex-position` | +| `batch_size` | `256` img / `64` video | both | Page size used when fetching items | +| `ffmpeg_path` | `ffmpeg` | videos | Path to the ffmpeg binary | +| `auto_download` | `true` | both | Download a missing model from Hugging Face | +| `hf_token` | none | both | Hugging Face token for gated models | +| `debug_log` | off | both | Write per-tag confidences to this JSON report (or `true` for `tagai-debug.json`) | +| `debug_floor` | `0.05` | both | Lowest score recorded in the debug report | +| `debug_top` | `10` | both | Candidates recorded per frame and model in the debug report | +| `markers` | `off` | videos | Create scene markers at the moments a tag fired: `off`, `ranges` or `best` | +| `marker_pad` | sampling interval | videos | Seconds padded onto each end of a marker (also the merge tolerance) | +| `marker_max` | `50` | videos | Maximum markers created per scene (0 = none) | +| `marker_gap` | `0` | videos | Merge markers of one tag closer than `2 * marker_pad + this` seconds | +| `marker_whitelist`| none | videos | Only these tags may create markers (comma-separated names) | +| `marker_blacklist`| none | videos | These tags never create markers (comma-separated names) | +| `tag_whitelist` | none | both | Only these tags are written to the item (comma-separated names) | +| `tag_blacklist` | none | both | These tags are never written to the item (comma-separated names) | + +## Notes + +* The plugin never writes anything but (optionally) JSON to stdout, since Stash + parses stdout. All logging goes through `stashapi.log`. +* Missing models produce a clear error pointing at + [`Models/README.md`](Models/README.md) rather than crashing the whole job. + +## Project layout + +Stash runs `tagAI.py` directly (`exec: [python, "{pluginDir}/tagAI.py"]`), so that +file has to stay put -- but it is only a short launcher for the `tagai_core` +package next to it: + +| Module | Responsibility | +| ------------------------ | --------------------------------------------------------- | +| `tagAI.py` | entry point: check dependencies, call `tagai_core.cli.main` | +| `tagai_core/config.py` | constants, the model registry and class-name tables | +| `tagai_core/stashlog.py` | logging (`stashapi.log`, falling back to stderr) | +| `tagai_core/stash_client.py` | GraphQL connection, tag lookup and creation | +| `tagai_core/models.py` | model sessions, Hugging Face download, label loading | +| `tagai_core/preprocess.py` | per-family image preprocessing | +| `tagai_core/postprocess.py` | per-family score post-processing and aggregation | +| `tagai_core/debug.py` | calibration debug report | +| `tagai_core/inference.py` | generic inference shared by the image and video pipelines | +| `tagai_core/images.py` | image queries and tag application | +| `tagai_core/frames.py` | video frame timing and ffmpeg extraction | +| `tagai_core/markers.py` | scene marker planning and creation | +| `tagai_core/scenes.py` | scene queries and per-scene processing | +| `tagai_core/jobs.py` | pagination loop and image/scene runners | +| `tagai_core/cli.py` | argument parsing, task dispatch (`main`) | + +The package imports in the order above (roughly leaf to top), so no module +imports a module that imports it back. Tests live in `tests/test_tagai.py` and +run with `python -m unittest discover -s tests -v`; they import the submodules +directly (e.g. `from tagai_core import cli`) and patch names on the module that +defines them. diff --git a/plugins/tagAI/_api.key b/plugins/tagAI/_api.key new file mode 100755 index 00000000..9368af7c --- /dev/null +++ b/plugins/tagAI/_api.key @@ -0,0 +1 @@ +paste api Key in here and rename file to api.key \ No newline at end of file diff --git a/plugins/tagAI/requirements.txt b/plugins/tagAI/requirements.txt new file mode 100755 index 00000000..85b563fc --- /dev/null +++ b/plugins/tagAI/requirements.txt @@ -0,0 +1,13 @@ +# TagAI runtime dependencies. +# +# ffmpeg is ALSO required for video tagging but is a system dependency, not a +# pip package. Make sure `ffmpeg` is on your PATH (or pass the ffmpeg_path arg). +# +# stashapp-tools provides `stashapi` (structured logging to the Stash UI). The +# plugin degrades gracefully to stderr logging when it is missing, but Stash +# itself bundles it, so it is normally available already. +onnxruntime-gpu +numpy +pillow +requests +stashapp-tools diff --git a/plugins/tagAI/tagAI.py b/plugins/tagAI/tagAI.py new file mode 100755 index 00000000..92dab05d --- /dev/null +++ b/plugins/tagAI/tagAI.py @@ -0,0 +1,33 @@ +"""TagAI Stash plugin entry point. + +Stash runs this file directly (see tagAI.yml, ``exec: [python, "{pluginDir}/tagAI.py"]``). +The implementation lives in the ``tagai_core`` package next to this file; this +module only bootstraps it. + +Stash passes the job arguments as JSON on stdin; this process must never write +anything but (optionally) JSON to stdout -- Stash parses stdout. All logging goes +through :mod:`tagai_core.stashlog` (stashapi.log, falling back to stderr). + +Package layout:: + + config constants, model registry, class-name tables + stashlog logging (stashapi.log + stderr fallback) + stash_client GraphQL connection, tag lookup/creation + models model sessions, Hugging Face download, label loading + preprocess per-family image preprocessing + postprocess per-family score post-processing and aggregation + debug calibration debug report + inference generic inference shared by the image and video pipelines + images image queries and tag application + frames video frame timing and ffmpeg extraction + markers scene marker planning and creation + scenes scene queries and per-scene processing + jobs pagination loop and image/scene runners + cli argument parsing and main() +""" + +from tagai_core.cli import main + + +if __name__ == "__main__": + main() diff --git a/plugins/tagAI/tagAI.yml b/plugins/tagAI/tagAI.yml new file mode 100755 index 00000000..12e82ae1 --- /dev/null +++ b/plugins/tagAI/tagAI.yml @@ -0,0 +1,125 @@ +name: TagAI +version: 1.8 +exec: + - python + - "{pluginDir}/tagAI.py" +description: Tag images and videos with locally hosted AI models +interface: raw +tasks: + - name: "Tag All Untagged Images" + description: "Tag all untagged images (WD tagger)" + defaultArgs: + target: images + model: wd-v1-4-moat-tagger-v2 + mode: TagAll + flag_tag: "AI_TagMe" + add_tag: "AI_Tagged" + tag_prefix: "AI: " + debug_log: "" + debug_floor: "0.05" + debug_top: "10" + tag_whitelist: "" + tag_blacklist: "" + - name: "Tag All Images -- Force" + description: "Tag all images, ignoring AI tag (WD tagger)" + defaultArgs: + target: images + model: wd-v1-4-moat-tagger-v2 + mode: ForceTag + flag_tag: "AI_TagMe" + add_tag: "AI_Tagged" + tag_prefix: "AI: " + debug_log: "" + debug_floor: "0.05" + debug_top: "10" + tag_whitelist: "" + tag_blacklist: "" + - name: "Tag All Flagged Images" + description: "Tag all images flagged with AI_TagMe (WD tagger)" + defaultArgs: + target: images + model: wd-v1-4-moat-tagger-v2 + mode: Tagged + flag_tag: "AI_TagMe" + add_tag: "AI_Tagged" + tag_prefix: "AI: " + debug_log: "" + debug_floor: "0.05" + debug_top: "10" + tag_whitelist: "" + tag_blacklist: "" + - name: "Tag All Untagged Videos" + description: "Tag all untagged videos (YOLO classifier)" + defaultArgs: + target: videos + model: "yolo11-nsfw,sex-position,wd-v1-4-moat-tagger-v2" + mode: TagAll + frame_interval: "5" + det_conf: "0.25" + tag_threshold: "0.65" + yolo_min_votes: "2" + sexpos_margin: "0.15" + flag_tag: "AI_TagMe" + add_tag: "AI_Tagged" + tag_prefix: "AI: " + debug_log: "" + debug_floor: "0.05" + debug_top: "10" + markers: "" + marker_pad: "5" + marker_max: "200" + marker_gap: "60" + marker_whitelist: "anal, blowjob, cowgirl position, doggystyle, fingering, girl on top, group sex, handjob, hardcore, imminent penetration, imminent vaginal, kissing penis, licking, licking penis, oral, pussy-licking, reverse cowgirl position, reverse upright straddle, sex, sex from behind, straddling, threesome, titjob, vaginal" + marker_blacklist: "" + tag_whitelist: "" + tag_blacklist: "age difference, artist name, blurry, blurry background, censored, closed eyes, food, food focus, fruit, loli, mosaic censoring, open pants, open fly, other, plant, realistic, rose, sketch, strawberry, uncensored, yaoi, yuri" + - name: "Tag All Videos -- Force" + description: "Tag all videos, ignoring AI tag (YOLO classifier)" + defaultArgs: + target: videos + model: "yolo11-nsfw,sex-position,wd-v1-4-moat-tagger-v2" + mode: ForceTag + frame_interval: "5" + det_conf: "0.25" + tag_threshold: "0.65" + yolo_min_votes: "2" + sexpos_margin: "0.15" + flag_tag: "AI_TagMe" + add_tag: "AI_Tagged" + tag_prefix: "AI: " + debug_log: "" + debug_floor: "0.05" + debug_top: "10" + markers: "" + marker_pad: "5" + marker_max: "200" + marker_gap: "60" + marker_whitelist: "anal, blowjob, cowgirl position, doggystyle, fingering, girl on top, group sex, handjob, hardcore, imminent penetration, imminent vaginal, kissing penis, licking, licking penis, oral, pussy-licking, reverse cowgirl position, reverse upright straddle, sex, sex from behind, straddling, threesome, titjob, vaginal" + marker_blacklist: "" + tag_whitelist: "" + tag_blacklist: "age difference, artist name, blurry, blurry background, censored, closed eyes, food, food focus, fruit, loli, mosaic censoring, open pants, open fly, other, plant, realistic, rose, sketch, strawberry, uncensored, yaoi, yuri" + - name: "Tag All Flagged Videos" + description: "Tag all videos flagged with AI_TagMe (YOLO classifier)" + defaultArgs: + target: videos + model: "yolo11-nsfw,sex-position,wd-v1-4-moat-tagger-v2" + mode: Tagged + frame_interval: "5" + det_conf: "0.25" + tag_threshold: "0.65" + yolo_min_votes: "2" + sexpos_margin: "0.15" + flag_tag: "AI_TagMe" + add_tag: "AI_Tagged" + tag_prefix: "AI: " + debug_log: "" + debug_floor: "0.05" + debug_top: "10" + markers: "ranges" + marker_pad: "5" + marker_max: "200" + marker_gap: "60" + marker_whitelist: "anal, blowjob, cowgirl position, doggystyle, fingering, girl on top, group sex, handjob, hardcore, imminent penetration, imminent vaginal, kissing penis, licking, licking penis, oral, pussy-licking, reverse cowgirl position, reverse upright straddle, sex, sex from behind, straddling, threesome, titjob, vaginal" + marker_blacklist: "" + tag_whitelist: "" + tag_blacklist: "age difference, artist name, blurry, blurry background, censored, closed eyes, food, food focus, fruit, loli, mosaic censoring, open pants, open fly, other, plant, realistic, rose, sketch, strawberry, uncensored, yaoi, yuri" diff --git a/plugins/tagAI/tagai_core/__init__.py b/plugins/tagAI/tagai_core/__init__.py new file mode 100755 index 00000000..e69de29b diff --git a/plugins/tagAI/tagai_core/cli.py b/plugins/tagAI/tagai_core/cli.py new file mode 100755 index 00000000..47a40f96 --- /dev/null +++ b/plugins/tagAI/tagai_core/cli.py @@ -0,0 +1,229 @@ +"""Argument parsing and the plugin entry point.""" + +from tagai_core import stash_client, stashlog +from tagai_core.debug import debug_start, resolve_debug_path +from tagai_core.images import get_images, get_images_force, get_images_w_tag +from tagai_core.jobs import run_image_job, run_scene_job +from tagai_core.markers import parse_marker_mode +from tagai_core.models import ensure_model_files, get_model +from tagai_core.scenes import get_scenes, get_scenes_force, get_scenes_w_tag +from tagai_core.stash_client import load_tag_cache, setup_connection +from tagai_core.tags import parse_tag_filter +from tagai_core.config import ( + CHAR_THRESHOLD, DEBUG_FLOOR_DEFAULT, DEBUG_TOP_DEFAULT, DEFAULT_IMAGE_MODEL, + DEFAULT_VIDEO_MODEL, FRAME_COUNT_DEFAULT, FRAME_INTERVAL_DEFAULT, + FRAME_WORKERS_DEFAULT, INFER_BATCH_SIZE_DEFAULT, JOYTAG_THRESHOLD_DEFAULT, + MARKER_MAX_DEFAULT, MAX_FRAMES_DEFAULT, PER_PAGE_IMAGE_DEFAULT, + PER_PAGE_VIDEO_DEFAULT, SEXPOS_MARGIN_DEFAULT, SEXPOS_THRESHOLD_DEFAULT, + TAG_THRESHOLD, YOLO_DET_CONF_DEFAULT, YOLO_DET_IOU_DEFAULT, YOLO_MARGIN_DEFAULT, + YOLO_MIN_VOTES_DEFAULT, YOLO_SHARE_DEFAULT, YOLO_THRESHOLD_DEFAULT, + YOLO_TOP_K_DEFAULT, _DEP_IMPORT_ERROR, +) +import json +import os +import sys + + +def parse_model_names(value, default): + """Normalise the ``model`` argument into an ordered list of model names. + + Accepts a single name or several separated by commas, semicolons or + whitespace (e.g. ``yolo11-nsfw,sex-position``). Duplicates are dropped and + an empty value falls back to ``default``. + """ + if isinstance(value, (list, tuple)): + raw = [str(item) for item in value] + else: + raw = str(value or "").replace(";", ",").replace(",", " ").split() + + names = [] + for item in raw: + item = item.strip() + if item and item not in names: + names.append(item) + if not names and default: + names = [default] + return names + + +def build_cfg(args): + """Normalise the runtime configuration from the plugin args.""" + + def as_int(key, default): + try: + return int(args.get(key, default)) + except (TypeError, ValueError): + return default + + def as_float(key, default): + try: + return float(args.get(key, default)) + except (TypeError, ValueError): + return default + + def as_bool(key, default=False): + v = args.get(key, default) + if isinstance(v, bool): + return v + if isinstance(v, str): + return v.strip().lower() in ("1", "true", "yes", "on") + return bool(v) + + return { + "tag_threshold": as_float("tag_threshold", TAG_THRESHOLD), + "char_threshold": as_float("char_threshold", CHAR_THRESHOLD), + "tag_characters": as_bool("tag_characters", False), + "joytag_threshold": as_float("joytag_threshold", JOYTAG_THRESHOLD_DEFAULT), + "det_conf": as_float("det_conf", YOLO_DET_CONF_DEFAULT), + "det_iou": as_float("det_iou", YOLO_DET_IOU_DEFAULT), + "sexpos_threshold": as_float("sexpos_threshold", SEXPOS_THRESHOLD_DEFAULT), + "sexpos_margin": as_float("sexpos_margin", SEXPOS_MARGIN_DEFAULT), + "yolo_threshold": as_float("yolo_threshold", YOLO_THRESHOLD_DEFAULT), + "yolo_margin": as_float("yolo_margin", YOLO_MARGIN_DEFAULT), + "frame_count": as_int("frame_count", FRAME_COUNT_DEFAULT), + "frame_interval": as_float("frame_interval", FRAME_INTERVAL_DEFAULT), + "max_frames": as_int("max_frames", MAX_FRAMES_DEFAULT), + "frame_workers": as_int("frame_workers", FRAME_WORKERS_DEFAULT), + "infer_batch_size": as_int("infer_batch_size", INFER_BATCH_SIZE_DEFAULT), + "yolo_multi": as_bool("yolo_multi", True), + "yolo_min_votes": as_int("yolo_min_votes", YOLO_MIN_VOTES_DEFAULT), + "yolo_share": as_float("yolo_share", YOLO_SHARE_DEFAULT), + "yolo_top_k": as_int("yolo_top_k", YOLO_TOP_K_DEFAULT), + "ffmpeg_path": args.get("ffmpeg_path") + or os.environ.get("TAGAI_FFMPEG") + or "ffmpeg", + "auto_download": as_bool("auto_download", True), + "debug_log": args.get("debug_log"), + "debug_floor": as_float("debug_floor", DEBUG_FLOOR_DEFAULT), + "debug_top": as_int("debug_top", DEBUG_TOP_DEFAULT), + "markers": parse_marker_mode(args.get("markers")), + "marker_pad": as_float("marker_pad", None), + "marker_max": as_int("marker_max", MARKER_MAX_DEFAULT), + "marker_gap": as_float("marker_gap", 0.0), + "marker_whitelist": parse_tag_filter( + args.get("marker_whitelist"), args.get("tag_prefix", "AI: ") + ), + "marker_blacklist": parse_tag_filter( + args.get("marker_blacklist"), args.get("tag_prefix", "AI: ") + ), + "tag_whitelist": parse_tag_filter( + args.get("tag_whitelist"), args.get("tag_prefix", "AI: ") + ), + "tag_blacklist": parse_tag_filter( + args.get("tag_blacklist"), args.get("tag_prefix", "AI: ") + ), + "hf_token": args.get("hf_token") or None, + } + + +def main(): + if _DEP_IMPORT_ERROR: + stashlog.log.error( + "TagAI is missing a required Python dependency: " + f"{_DEP_IMPORT_ERROR}. Install the requirements into the Python " + "interpreter Stash uses for plugins (see requirements.txt), e.g. " + "'pip install -r requirements.txt'." + ) + return + + raw = sys.stdin.read() + try: + json_input = json.loads(raw) if raw.strip() else {} + except ValueError as exc: + stashlog.log.error(f"Invalid JSON on stdin: {exc}") + json_input = {} + + setup_connection(json_input) + + args = (json_input or {}).get("args") or {} + + target = (args.get("target") or "images").lower() + if target in ("video", "scene", "scenes"): + target = "videos" + if target not in ("images", "videos"): + stashlog.log.error(f"Unknown target {target!r}; expected 'images' or 'videos'") + return + + if target == "videos": + default_model = DEFAULT_VIDEO_MODEL + default_per_page = PER_PAGE_VIDEO_DEFAULT + else: + default_model = DEFAULT_IMAGE_MODEL + default_per_page = PER_PAGE_IMAGE_DEFAULT + + model_names = parse_model_names(args.get("model"), default_model) + + mode = args.get("mode") or "TagAll" + if mode not in ("TagAll", "ForceTag", "Tagged"): + stashlog.log.error( + f"Unknown mode {mode!r}; expected TagAll, ForceTag or Tagged" + ) + return + + try: + per_page = int(args.get("batch_size") or default_per_page) + except (TypeError, ValueError): + per_page = default_per_page + + prefix = args.get("tag_prefix", "AI: ") + cfg = build_cfg(args) + cfg["debug"] = debug_start(resolve_debug_path(cfg.get("debug_log")), cfg) + + load_tag_cache() + flag_tag_id = stash_client.ensure_tag(args.get("flag_tag") or "AI_TagMe") + ai_tagged_id = stash_client.ensure_tag(args.get("add_tag") or "AI_Tagged") + if not flag_tag_id or not ai_tagged_id: + stashlog.log.error("TagAI could not create the required flag/tagged tags; aborting") + return + + models = [] + for name in model_names: + try: + ensure_model_files( + name, + token=cfg.get("hf_token"), + allow_download=cfg.get("auto_download", True), + on_progress=lambda p: stashlog.log.progress(p), + ) + except Exception as exc: + stashlog.log.error(f"Could not obtain model {name!r}: {exc}") + continue + + try: + models.append(get_model(name)) + except Exception as exc: + stashlog.log.error(f"Could not load model {name!r}: {exc}") + + if not models: + stashlog.log.error("TagAI could not load any model; aborting") + return + + stashlog.log.info( + f"TagAI: target={target} models={','.join(m['name'] for m in models)} " + f"mode={mode} per_page={per_page}" + ) + + if target == "images": + if mode == "TagAll": + fetch = lambda page: get_images(ai_tagged_id, page, per_page) + elif mode == "ForceTag": + fetch = lambda page: get_images_force(page, per_page) + else: + fetch = lambda page: get_images_w_tag(flag_tag_id, page, per_page) + + processed = run_image_job( + fetch, mode, per_page, models, cfg, flag_tag_id, ai_tagged_id, prefix + ) + else: + if mode == "TagAll": + fetch = lambda page: get_scenes(ai_tagged_id, page, per_page) + elif mode == "ForceTag": + fetch = lambda page: get_scenes_force(page, per_page) + else: + fetch = lambda page: get_scenes_w_tag(flag_tag_id, page, per_page) + + processed = run_scene_job( + fetch, mode, per_page, models, cfg, flag_tag_id, ai_tagged_id, prefix + ) + + stashlog.log.info(f"TagAI finished: processed {len(processed)} item(s)") diff --git a/plugins/tagAI/tagai_core/config.py b/plugins/tagAI/tagai_core/config.py new file mode 100755 index 00000000..05db4b65 --- /dev/null +++ b/plugins/tagAI/tagai_core/config.py @@ -0,0 +1,286 @@ +"""Configuration constants, model registry and class-name tables.""" + +import os + + +_DEP_IMPORT_ERROR = None +try: + import numpy as np + import requests + from PIL import Image +except ImportError as _dep_exc: # pragma: no cover - only hit when deps missing + _DEP_IMPORT_ERROR = str(_dep_exc) + np = None + requests = None + Image = None + + +# ========================================================================== # +# 1. CONFIG +# ========================================================================== # + +PLUGIN_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +MODEL_DIR = os.path.join(PLUGIN_DIR, "Models") +API_KEY_PATH = os.path.join(PLUGIN_DIR, "api.key") + +STASH_URL_DEFAULT = "http://localhost:9999/graphql" + +TAG_THRESHOLD = 0.5 # WD general tag threshold +CHAR_THRESHOLD = 0.85 # WD character tag threshold +BATCH_LIMIT = 1024 +PER_PAGE_IMAGE_DEFAULT = 256 +PER_PAGE_VIDEO_DEFAULT = 64 +FRAME_COUNT_DEFAULT = 8 +FRAME_INTERVAL_DEFAULT = 10.0 # seconds between sampled frames (0 = use frame_count) +MAX_FRAMES_DEFAULT = 512 # upper bound on frames sampled per video +FRAME_LONG_SIDE = 448 # intermediate frame extraction size (long side) +FRAME_WORKERS_DEFAULT = 4 # parallel ffmpeg processes used to extract frames +INFER_BATCH_SIZE_DEFAULT = 8 # frames per onnxruntime call (video pipeline) +YOLO_THRESHOLD_DEFAULT = 0.5 +YOLO_MARGIN_DEFAULT = 0.4 +YOLO_MIN_VOTES_DEFAULT = 1 # min accepted frames for a YOLO class to be tagged +YOLO_SHARE_DEFAULT = 0.0 # min fraction of accepted frames (0 = no share filter) +YOLO_TOP_K_DEFAULT = 1 # classes accepted per frame (1 = top-1 only) +JOYTAG_SIZE = 448 # JoyTag expects a 448x448 square +JOYTAG_THRESHOLD_DEFAULT = 0.4 # JoyTag sigmoid threshold (author's default) + +# YOLO11-NSFW object detector (explicit acts / body parts). +YOLO_DET_SIZE = 640 # detection input size +YOLO_DET_CONF_DEFAULT = 0.25 # detection confidence threshold +YOLO_DET_IOU_DEFAULT = 0.45 # NMS IoU threshold +# Used when the ONNX metadata does not carry the class names. +YOLO_DET_CLASS_NAMES = [ + "person", "breast", "vulva", "butt", "male", + "penis", "anal", "vaginal", "blowjob", "handjob", +] + +# porntech/sex-position: single-label sex-position classifier (resnet50). +SEXPOS_SIZE = 224 # centre-cropped input size +SEXPOS_RESIZE = 256 # shorter side before the centre crop +SEXPOS_THRESHOLD_DEFAULT = 0.5 # top-1 confidence needed to tag the frame +SEXPOS_MARGIN_DEFAULT = 0.0 # optional top1-top2 gap (0 = disabled) +SEXPOS_CLASS_NAMES = [ + "blowjob", "hardcore", "titjob", "handjob", + "pussy-licking", "fingering", "other", "solo", +] +# Debug report: per-tag confidences of one run, written to a JSON file so the +# thresholds can be calibrated against a known video (see README). +DEBUG_FILENAME = "tagai-debug.json" +DEBUG_FLOOR_DEFAULT = 0.05 # candidates below this score are not recorded +DEBUG_TOP_DEFAULT = 10 # candidates recorded per frame and model +DEBUG_AGG_TOP_DEFAULT = 100 # candidates recorded in the per-model summary + +# Scene markers: optional per-tag markers at the moments a tag fired. The +# marker spans are padded by the sampling interval so a marker always lasts at +# least ``2 * marker_pad`` seconds (see build_markers). +MARKER_MODES = ("off", "ranges", "best") +MARKER_MAX_DEFAULT = 50 # upper bound on markers created per scene +MARKER_MIN_LENGTH = 1.0 # shortest allowed marker span (seconds) + +# ImageNet statistics used by the timm preprocessing of that model. +DEFAULT_IMAGE_MODEL = "wd-v1-4-moat-tagger-v2" +DEFAULT_VIDEO_MODEL = "yolo26n-cls-porn" + +# Hugging Face download settings (used to fetch missing models automatically). +HF_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co").rstrip("/") +HF_TOKEN_ENV_VARS = ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_TOKEN") +HF_TOKEN_FILENAMES = ("hf_token", "hf.token", "huggingface.token") + +# Model registry. ``family`` selects the strategy layer. +# Model registry. ``family`` selects the strategy layer; ``repo``/``files`` +# describe how to fetch the model from Hugging Face when it is not on disk yet. +# ``files`` is a list of (local_filename, path_in_repo). +MODELS = { + "wd-v1-4-moat-tagger-v2": { + "family": "wd", + "onnx": "wd-v1-4-moat-tagger-v2.onnx", + "labels": "wd-v1-4-moat-tagger-v2.csv", + "repo": "SmilingWolf/wd-v1-4-moat-tagger-v2", + "files": [ + ("wd-v1-4-moat-tagger-v2.onnx", "model.onnx"), + ("wd-v1-4-moat-tagger-v2.csv", "selected_tags.csv"), + ], + }, + "wd-convnext-tagger-v3": { + "family": "wd", + "onnx": "wd-convnext-tagger-v3.onnx", + "labels": "wd-convnext-tagger-v3.csv", + "repo": "SmilingWolf/wd-convnext-tagger-v3", + "files": [ + ("wd-convnext-tagger-v3.onnx", "model.onnx"), + ("wd-convnext-tagger-v3.csv", "selected_tags.csv"), + ], + }, + "wd-v1-4-vit-tagger": { + "family": "wd", + "onnx": "wd-v1-4-vit-tagger.onnx", + "labels": "wd-v1-4-vit-tagger.csv", + "repo": "SmilingWolf/wd-v1-4-vit-tagger", + "files": [ + ("wd-v1-4-vit-tagger.onnx", "model.onnx"), + ("wd-v1-4-vit-tagger.csv", "selected_tags.csv"), + ], + }, + # WD v3 taggers: 10,861 tags (vs 9,083 in the v1-4 series) and, unlike the + # v1-4 graphs, a dynamic batch dimension. + "wd-swinv2-tagger-v3": { + "family": "wd", + "onnx": "wd-swinv2-tagger-v3.onnx", + "labels": "wd-swinv2-tagger-v3.csv", + "repo": "SmilingWolf/wd-swinv2-tagger-v3", + "files": [ + ("wd-swinv2-tagger-v3.onnx", "model.onnx"), + ("wd-swinv2-tagger-v3.csv", "selected_tags.csv"), + ], + }, + "wd-eva02-large-tagger-v3": { + "family": "wd", + "onnx": "wd-eva02-large-tagger-v3.onnx", + "labels": "wd-eva02-large-tagger-v3.csv", + "repo": "SmilingWolf/wd-eva02-large-tagger-v3", + "files": [ + ("wd-eva02-large-tagger-v3.onnx", "model.onnx"), + ("wd-eva02-large-tagger-v3.csv", "selected_tags.csv"), + ], + }, + # JoyTag: ViT-B/16 multi-label tagger trained on Danbooru plus manually + # tagged photographic images, so it works well on real (photo) porn. + "joytag": { + "family": "joytag", + "onnx": "joytag.onnx", + "labels": "joytag-top_tags.txt", + "repo": "fancyfeast/joytag", + "files": [ + ("joytag.onnx", "model.onnx"), + ("joytag-top_tags.txt", "top_tags.txt"), + ], + }, + # YOLO11-NSFW: object detection of explicit acts and body parts (10 + # classes). The upstream repository only ships a .pt checkpoint, so the + # ONNX has to be exported once by hand (see Models/README.md) -- it cannot + # be downloaded automatically. + "yolo11-nsfw": { + "family": "yolo-det", + "onnx": "yolo11-nsfw.onnx", + }, + # porntech/sex-position: single-label sex-position classifier (resnet50, 8 + # classes). The upstream repository ships only pytorch_model.bin, so the + # ONNX has to be exported once by hand (see Models/README.md) -- it cannot + # be downloaded automatically. + "sex-position": { + "family": "sexpos", + "onnx": "sex-position.onnx", + }, + "yolo26n-cls-porn": { + "family": "yolo", + "onnx": "yolo26n-cls-porn.onnx", + "repo": "likewendy/yolo26n-cls-porn", + "gated": True, + "files": [ + ("yolo26n-cls-porn.onnx", "weights/best.onnx"), + ], + # Optional: per-class confidence thresholds published by the author. + # Fetched best-effort; the bundled defaults are used when absent. + "optional_files": [ + ("yolo26n-cls-porn-thresholds.json", + "class35_conf_thresholds_by_name.json"), + ], + }, +} + +# The YOLO classifier (likewendy/yolo26n-cls-porn) ships Chinese class names in +# its ONNX metadata. They are translated to English so Stash tags stay +# readable. The order of ``YOLO_CLASS_NAMES`` below is the model's real output +# order (verified against the ONNX metadata), NOT the order the README displays. +YOLO_ZH_TO_EN = { + "AI生成": "AI Generated", + "乳交": "Titfuck", + "偷窥": "Voyeurism", + "公共暴露": "Public Exposure", + "内射": "Creampie", + "动漫": "Anime", + "卷发": "Curly Hair", + "双插": "Double Penetration", + "口交": "Blowjob", + "口爆大会": "Bukkake", + "圣诞": "Christmas", + "女仆": "Maid", + "女女假阳具": "Lesbian Strapon", + "女女摩擦": "Lesbian Tribbing", + "女着男裸": "CFNM", + "孕妇": "Pregnant", + "恋尿": "Pissing", + "手交": "Handjob", + "接吻": "Kissing", + "油亮": "Oiled", + "潮吹": "Squirting", + "男性独处": "Solo Male", + "男性群像": "Male Group", + "皮革乳胶": "Leather and Latex", + "纹身": "Tattoo", + "绿帽视角": "Cuckold", + "肛门特写": "Asshole Close-up", + "脚交": "Footjob", + "腋窝": "Armpit", + "舔肛": "Rimjob", + "舔阴": "Pussy Licking", + "裤装": "Pants", + "跨种族": "Interracial", + "镜前": "Mirror", + "骑脸": "Facesitting", +} + +# Fallback English class names in the model's output order (used only when the +# ONNX metadata does not carry names). +YOLO_CLASS_NAMES = [ + "AI Generated", "Titfuck", "Voyeurism", "Public Exposure", "Creampie", + "Anime", "Curly Hair", "Double Penetration", "Blowjob", "Bukkake", + "Christmas", "Maid", "Lesbian Strapon", "Lesbian Tribbing", "CFNM", + "Pregnant", "Pissing", "Handjob", "Kissing", "Oiled", + "Squirting", "Solo Male", "Male Group", "Leather and Latex", "Tattoo", + "Cuckold", "Asshole Close-up", "Footjob", "Armpit", "Rimjob", + "Pussy Licking", "Pants", "Interracial", "Mirror", "Facesitting", +] + +# Per-class confidence thresholds published by the model author (converted from +# class35_conf_thresholds_by_name.json). Using these instead of a flat 0.5 cuts +# down on false positives such as "Kissing" / "Pants". +YOLO_DEFAULT_THRESHOLDS = { + "AI Generated": 0.5, + "Titfuck": 0.5, + "Voyeurism": 0.51236963, + "Public Exposure": 0.55655706, + "Creampie": 0.71145976, + "Anime": 0.5, + "Curly Hair": 0.74997914, + "Double Penetration": 0.5, + "Blowjob": 0.89571738, + "Bukkake": 0.5, + "Christmas": 0.5, + "Maid": 0.60505283, + "Lesbian Strapon": 0.5, + "Lesbian Tribbing": 0.51000375, + "CFNM": 0.66284251, + "Pregnant": 0.51583439, + "Pissing": 0.58715159, + "Handjob": 0.88127387, + "Kissing": 0.64717263, + "Oiled": 0.85343647, + "Squirting": 0.85944873, + "Solo Male": 0.84280539, + "Male Group": 0.76054007, + "Leather and Latex": 0.56601757, + "Tattoo": 0.80576283, + "Cuckold": 0.9, + "Asshole Close-up": 0.5, + "Footjob": 0.5, + "Armpit": 0.5, + "Rimjob": 0.67189211, + "Pussy Licking": 0.70615655, + "Pants": 0.96068484, + "Interracial": 0.8965643, + "Mirror": 0.5, + "Facesitting": 0.83216292, +} + +# Runtime connection state (populated from stdin / api.key by setup_connection). diff --git a/plugins/tagAI/tagai_core/debug.py b/plugins/tagAI/tagai_core/debug.py new file mode 100755 index 00000000..6d812506 --- /dev/null +++ b/plugins/tagAI/tagai_core/debug.py @@ -0,0 +1,130 @@ +"""Calibration debug report (JSON) helpers.""" + +from tagai_core import config, stashlog +from tagai_core.config import DEBUG_AGG_TOP_DEFAULT, DEBUG_FILENAME +import json +import os +import time + + +def aggregate_candidates(per_frame): + """Per-tag max of several candidate lists, best first.""" + best = {} + for candidates in per_frame: + for name, score in candidates: + if name not in best or score > best[name]: + best[name] = float(score) + return sorted(best.items(), key=lambda kv: kv[1], reverse=True) + + +def _debug_settings(model, cfg): + """The cfg values that decide whether this model's tags are accepted.""" + family = model["family"] + if family == "wd": + return { + "tag_threshold": cfg.get("tag_threshold"), + "char_threshold": cfg.get("char_threshold"), + "tag_characters": cfg.get("tag_characters"), + } + if family == "joytag": + return {"joytag_threshold": cfg.get("joytag_threshold")} + if family == "yolo-det": + return { + "det_conf": cfg.get("det_conf"), + "det_iou": cfg.get("det_iou"), + } + if family == "sexpos": + return { + "sexpos_threshold": cfg.get("sexpos_threshold"), + "sexpos_margin": cfg.get("sexpos_margin"), + } + return { + "yolo_threshold": cfg.get("yolo_threshold"), + "yolo_margin": cfg.get("yolo_margin"), + "yolo_top_k": cfg.get("yolo_top_k"), + "yolo_multi": cfg.get("yolo_multi"), + "yolo_min_votes": cfg.get("yolo_min_votes"), + "yolo_share": cfg.get("yolo_share"), + "thresholds": model.get("thresholds"), + } + + +def _debug_entry(model, cfg): + """Start a debug record for one model.""" + return { + "name": model.get("name"), + "family": model["family"], + "settings": _debug_settings(model, cfg), + "per_frame": [], + } + + +def _finish_debug_entry(entry, tags): + """Store the aggregated candidates and the accepted tags in an entry.""" + if entry is None: + return + accepted = [name for name, _score in tags] + entry["tags"] = accepted + aggregate = aggregate_candidates(entry.get("per_frame") or []) + entry["aggregate"] = [ + {"tag": name, "score": round(float(score), 4), "accepted": name in accepted} + for name, score in aggregate[:DEBUG_AGG_TOP_DEFAULT] + ] + + +def _debug_pairs(candidates): + """Compact ``[[tag, score], ...]`` rows for the debug report.""" + return [[name, round(float(score), 4)] for name, score in candidates] + + +_DEBUG = {"path": None, "items": [], "settings": {}} + + +def resolve_debug_path(value): + """Map the ``debug_log`` argument to a report path (``None`` = disabled).""" + if value is None or isinstance(value, bool): + return os.path.join(config.PLUGIN_DIR, DEBUG_FILENAME) if value else None + text = str(value).strip() + if not text or text.lower() in ("false", "0", "no", "off"): + return None + if text.lower() in ("true", "1", "yes", "on"): + return os.path.join(config.PLUGIN_DIR, DEBUG_FILENAME) + return text + + +def debug_start(path, cfg): + """Enable the debug report; returns whether it is active.""" + if not path: + _DEBUG["path"] = None + return False + _DEBUG["path"] = path + _DEBUG["items"] = [] + _DEBUG["settings"] = { + key: value for key, value in (cfg or {}).items() if key != "hf_token" + } + stashlog.log.info(f"TagAI debug: writing tag confidences to {path}") + return True + + +def debug_add_item(item): + """Append one item record to the report and rewrite the JSON file.""" + path = _DEBUG.get("path") + if not path: + return + _DEBUG["items"].append(item) + report = { + "plugin": "TagAI", + "format": 1, + "generated": time.strftime("%Y-%m-%dT%H:%M:%S"), + "settings": _DEBUG.get("settings") or {}, + "items": _DEBUG["items"], + } + try: + directory = os.path.dirname(os.path.abspath(path)) + if directory: + os.makedirs(directory, exist_ok=True) + with open(path, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=1) + stashlog.log.info(f"TagAI debug: {len(_DEBUG['items'])} item(s) in {path}") + except Exception as exc: + stashlog.log.error(f"TagAI debug: could not write {path}: {exc}") diff --git a/plugins/tagAI/tagai_core/frames.py b/plugins/tagAI/tagai_core/frames.py new file mode 100755 index 00000000..b5009082 --- /dev/null +++ b/plugins/tagAI/tagai_core/frames.py @@ -0,0 +1,124 @@ +"""Video frame time computation and ffmpeg frame extraction.""" + +from tagai_core import stashlog +from tagai_core.config import ( + FRAME_COUNT_DEFAULT, FRAME_INTERVAL_DEFAULT, FRAME_LONG_SIDE, FRAME_WORKERS_DEFAULT, + MAX_FRAMES_DEFAULT, +) +from concurrent.futures import ThreadPoolExecutor +import os +import subprocess + + +def compute_frame_times(duration, frame_count=FRAME_COUNT_DEFAULT, + frame_interval=FRAME_INTERVAL_DEFAULT, + max_frames=MAX_FRAMES_DEFAULT): + """Return sample timestamps (seconds) inside the 10%-90% duration band. + + ``frame_interval > 0`` samples one frame roughly every ``frame_interval`` + seconds; otherwise ``frame_count`` frames are spread evenly. The number of + samples is always capped at ``max_frames`` so long videos stay cheap. + """ + if not duration or duration <= 0: + return [] + + try: + max_frames = max(1, int(max_frames)) + except (TypeError, ValueError): + max_frames = MAX_FRAMES_DEFAULT + + lo = 0.1 * float(duration) + span = 0.8 * float(duration) + + if frame_interval and float(frame_interval) > 0: + n = int(span // float(frame_interval)) + 1 + else: + try: + n = int(frame_count) + except (TypeError, ValueError): + n = FRAME_COUNT_DEFAULT + n = n if n > 0 else 1 + + n = max(1, min(n, max_frames)) + return [lo + span * (i + 0.5) / n for i in range(n)] + + +def extract_frames(video_path, times, out_dir, ffmpeg_path="ffmpeg", + long_side=FRAME_LONG_SIDE, workers=FRAME_WORKERS_DEFAULT, + include_times=False): + """Extract one frame per timestamp in ``times`` (seconds). + + Frames are written at an aspect-preserving ``long_side`` without padding so + each model family can apply its own final transform. With ``workers > 1`` + the per-timestamp ffmpeg processes run concurrently, which mostly hides the + seek latency of slow/networked storage. + + With ``include_times`` the returned list holds ``(path, seconds)`` pairs so + callers can map a frame back to its position in the video (used to place + scene markers); the default keeps the historical list of paths. + """ + frames = [] + if not video_path or times is None: + return frames + times = list(times) + if not times: + return frames + if not os.path.isfile(video_path): + stashlog.log.error(f"Video file not found: {video_path}") + return frames + + try: + workers = max(1, int(workers)) + except (TypeError, ValueError): + workers = 1 + + vf = ( + f"scale='if(gt(iw,ih),{long_side},-2)'" + f":'if(gt(iw,ih),-2,{long_side})'" + ) + + def extract_one(index, t): + """Run one ffmpeg invocation; returns (ok, index, message).""" + t = float(t) + out = os.path.join(out_dir, f"frame_{index:03d}.jpg") + cmd = [ + ffmpeg_path, "-v", "error", + "-ss", f"{t:.3f}", + "-i", video_path, + "-vf", vf, + "-frames:v", "1", + "-y", out, + ] + try: + proc = subprocess.run(cmd, capture_output=True, text=True) + except FileNotFoundError: + return ("ffmpeg_missing", index, ffmpeg_path) + + if proc.returncode != 0 or not os.path.isfile(out): + return ("failed", index, f"t={t:.3f}s: {proc.stderr.strip()}") + return ("ok", index, out) + + jobs = list(enumerate(times)) + if workers > 1 and len(jobs) > 1: + with ThreadPoolExecutor(max_workers=workers) as pool: + outcomes = list(pool.map(lambda job: extract_one(*job), jobs)) + else: + outcomes = [extract_one(*job) for job in jobs] + + if any(status == "ffmpeg_missing" for status, _i, _m in outcomes): + stashlog.log.error( + f"ffmpeg not found ({ffmpeg_path}); cannot extract frames. " + "Install ffmpeg or pass the ffmpeg_path argument." + ) + return frames + + for status, index, message in sorted(outcomes, key=lambda item: item[1]): + if status == "ok": + if include_times: + frames.append((message, float(times[index]))) + else: + frames.append(message) + else: + stashlog.log.error(f"ffmpeg failed for {video_path} at {message}") + + return frames diff --git a/plugins/tagAI/tagai_core/images.py b/plugins/tagAI/tagai_core/images.py new file mode 100755 index 00000000..4b2800f6 --- /dev/null +++ b/plugins/tagAI/tagai_core/images.py @@ -0,0 +1,87 @@ +"""Image GraphQL queries and tag application.""" + +from tagai_core import stash_client, stashlog +from tagai_core.config import PER_PAGE_IMAGE_DEFAULT + + +def get_images(ai_tagged_id, page=1, per_page=PER_PAGE_IMAGE_DEFAULT): + """Untagged images (excluding those already carrying ``ai_tagged_id``).""" + q = """ +query MyQuery($excludes: [ID!], $page: Int, $per_page: Int) { + findImages( + image_filter: { tags: { modifier: INCLUDES_ALL, excludes: $excludes } } + filter: { page: $page, per_page: $per_page } + ) { + images { id files { path } tags { id } } + count + } +}""" + res = stash_client.gql(q, {"excludes": [ai_tagged_id], "page": page, "per_page": per_page}) + if "errors" in res: + return [], 0 + data = res.get("data", {}).get("findImages") or {} + return data.get("images") or [], data.get("count") or 0 + + +def get_images_force(page=1, per_page=PER_PAGE_IMAGE_DEFAULT): + q = """ +query MyQuery($page: Int, $per_page: Int) { + findImages(filter: { page: $page, per_page: $per_page }) { + images { id files { path } tags { id } } + count + } +}""" + res = stash_client.gql(q, {"page": page, "per_page": per_page}) + if "errors" in res: + return [], 0 + data = res.get("data", {}).get("findImages") or {} + return data.get("images") or [], data.get("count") or 0 + + +def get_images_w_tag(flag_tag_id, page=1, per_page=PER_PAGE_IMAGE_DEFAULT): + q = """ +query MyQuery($value: [ID!], $page: Int, $per_page: Int) { + findImages( + image_filter: { tags: { modifier: INCLUDES_ALL, value: $value } } + filter: { page: $page, per_page: $per_page } + ) { + images { id files { path } tags { id } } + count + } +}""" + res = stash_client.gql(q, {"value": [flag_tag_id], "page": page, "per_page": per_page}) + if "errors" in res: + return [], 0 + data = res.get("data", {}).get("findImages") or {} + return data.get("images") or [], data.get("count") or 0 + + +def apply_tags(img_id, new_tag_ids, exclude_tag_id=None, ai_tagged_id=None): + """Merge new tags onto an image, dropping ``exclude_tag_id`` and adding + ``ai_tagged_id``.""" + new_tag_ids = [t for t in new_tag_ids if t] + if ai_tagged_id: + new_tag_ids.append(ai_tagged_id) + if not new_tag_ids: + stashlog.log.info("No tags to apply.") + return + + q = """query($id:ID!){ findImage(id:$id){ tags{id} } }""" + r = stash_client.gql(q, {"id": img_id}) + existing = [ + t["id"] + for t in r.get("data", {}).get("findImage", {}).get("tags", []) + if t["id"] != exclude_tag_id + ] + + final = list(set(existing + new_tag_ids)) + + q = """mutation($id:ID!, $tags:[ID!]!){ + imageUpdate(input:{id:$id, tag_ids:$tags}){id} + }""" + stash_client.gql(q, {"id": img_id, "tags": final}) + + +# ========================================================================== # +# 5. VIDEO (SCENE) PIPELINE +# ========================================================================== # diff --git a/plugins/tagAI/tagai_core/inference.py b/plugins/tagAI/tagai_core/inference.py new file mode 100755 index 00000000..68b98e99 --- /dev/null +++ b/plugins/tagAI/tagai_core/inference.py @@ -0,0 +1,403 @@ +"""Generic model inference shared by the image and video pipelines.""" + +from tagai_core import postprocess, stashlog +from tagai_core.debug import _debug_entry, _debug_pairs, _finish_debug_entry +from tagai_core.postprocess import ( + _as_probabilities, _as_sigmoid_probabilities, _sexpos_from_probs, + joytag_postprocess, softmax, wd_aggregate, wd_postprocess, yolo_aggregate, + yolo_postprocess, +) +from tagai_core.preprocess import ( + joytag_preprocess, sexpos_preprocess, wd_preprocess, yolo_det_preprocess, + yolo_preprocess, yolo_preprocess_batch, +) +from tagai_core.config import ( + DEBUG_FLOOR_DEFAULT, DEBUG_TOP_DEFAULT, INFER_BATCH_SIZE_DEFAULT, + JOYTAG_THRESHOLD_DEFAULT, SEXPOS_MARGIN_DEFAULT, SEXPOS_THRESHOLD_DEFAULT, + YOLO_DET_CONF_DEFAULT, YOLO_DET_IOU_DEFAULT, YOLO_MIN_VOTES_DEFAULT, + YOLO_SHARE_DEFAULT, YOLO_TOP_K_DEFAULT, +) +from tagai_core.config import Image +from tagai_core.config import np +import zipfile + + +def load_pil_image(path): + """Open an image path, transparently supporting images inside .zip files.""" + if ".zip" in path.lower(): + zp, inner = path.split(".zip", 1) + zp += ".zip" + inner = inner.lstrip("\\/").replace("\\", "/") + with zipfile.ZipFile(zp, "r") as z: + fn = inner.split("/")[-1] + cand = [x for x in z.namelist() if x.endswith(fn)] + if not cand: + raise FileNotFoundError(fn) + return Image.open(z.open(cand[0])).convert("RGB") + return Image.open(path).convert("RGB") + + +def run_inference(img, model, cfg): + """Run one model on one image -> ``(family, raw scores)``. + + The returned scores are family specific: a probability vector for the + classifier families and the raw detection head for ``yolo-det`` (so the + debug report can re-decode it at a lower confidence floor). + """ + family = model["family"] + + if family == "wd": + arr = wd_preprocess(img, model["size"]) + out = model["session"].run( + [model["output_name"]], {model["input_name"]: arr} + )[0][0] + return family, out + + if family == "joytag": + arr = joytag_preprocess(img, model["size"]) + out = model["session"].run( + [model["output_name"]], {model["input_name"]: arr} + )[0][0] + return family, _as_sigmoid_probabilities(out) + + if family == "yolo-det": + arr, _meta = yolo_det_preprocess(img, model["size"]) + out = model["session"].run( + [model["output_name"]], {model["input_name"]: arr} + )[0] + return family, out + + if family == "sexpos": + arr = sexpos_preprocess(img, model["size"]) + out = model["session"].run( + [model["output_name"]], {model["input_name"]: arr} + )[0][0] + return family, softmax(np.asarray(out, dtype=np.float32).reshape(-1)) + + # yolo + arr = yolo_preprocess(img, model["size"]) + out = model["session"].run( + [model["output_name"]], {model["input_name"]: arr} + )[0][0] + return family, _as_probabilities(out) + + +def postprocess_scores(family, scores, model, cfg): + """Turn one model's raw scores into the accepted ``[(name, score)]``.""" + if family == "wd": + return wd_postprocess( + scores, + model, + tag_threshold=cfg["tag_threshold"], + char_threshold=cfg["char_threshold"], + tag_characters=cfg["tag_characters"], + ) + + if family == "joytag": + return joytag_postprocess( + scores, + model, + threshold=cfg.get("joytag_threshold", JOYTAG_THRESHOLD_DEFAULT), + ) + + if family == "yolo-det": + return postprocess.yolo_det_postprocess( + scores, + model, + conf_threshold=cfg.get("det_conf", YOLO_DET_CONF_DEFAULT), + iou_threshold=cfg.get("det_iou", YOLO_DET_IOU_DEFAULT), + ) + + if family == "sexpos": + return _sexpos_from_probs( + scores, + model, + threshold=cfg.get("sexpos_threshold", SEXPOS_THRESHOLD_DEFAULT), + margin=cfg.get("sexpos_margin", SEXPOS_MARGIN_DEFAULT), + ) + + return yolo_postprocess( + scores, + model, + thresholds=model.get("thresholds"), + margin=cfg["yolo_margin"], + default_threshold=cfg["yolo_threshold"], + top_k=cfg.get("yolo_top_k", YOLO_TOP_K_DEFAULT), + ) + + +def classify_image(img, model, cfg): + """Run the selected model on one PIL image -> ``[(name, score)]``.""" + family, scores = run_inference(img, model, cfg) + return postprocess_scores(family, scores, model, cfg) + + +def candidate_scores(family, scores, model, cfg, floor=None, top=None): + """Every score above ``floor``, best first (used by the debug report). + + Unlike :func:`postprocess_scores` this ignores the acceptance thresholds, + so the report shows how close the rejected tags were to being tagged. + """ + if floor is None: + floor = cfg.get("debug_floor", DEBUG_FLOOR_DEFAULT) + if top is None: + top = cfg.get("debug_top", DEBUG_TOP_DEFAULT) + + if family == "yolo-det": + detections = postprocess.yolo_det_postprocess( + scores, + model, + conf_threshold=floor, + iou_threshold=cfg.get("det_iou", YOLO_DET_IOU_DEFAULT), + ) + return [(name, float(score)) for name, score in detections[:top]] + + arr = np.asarray(scores, dtype=np.float32).reshape(-1) + if arr.size == 0: + return [] + + names = model.get("tags") or model.get("class_names") or [] + if len(names) != arr.size: + names = [f"class_{i}" for i in range(arr.size)] + + hits = [] + for idx in np.argsort(arr)[::-1]: + idx = int(idx) + score = float(arr[idx]) + if score < floor: + break + hits.append((names[idx], score)) + if len(hits) >= max(1, int(top)): + break + return hits + + +def merge_tag_results(results): + """Merge several ``[(name, score)]`` lists, keeping the best score.""" + best = {} + for res in results: + for name, score in res: + if name not in best or score > best[name]: + best[name] = float(score) + return sorted(best.items(), key=lambda kv: kv[1], reverse=True) + + +def classify_with_models(img, models, cfg, debug=None): + """Run every model on one image and merge the tags they produce. + + ``models`` may be a single model entry (dict) or a list of them. With one + model the result is returned unchanged, so the tag order of the existing + single-model pipelines is preserved exactly. + + When ``debug`` is a list, one record per model (with the per-tag + confidences, including the rejected ones) is appended to it. + """ + if isinstance(models, dict): + models = [models] + results = [] + for model in models: + try: + if debug is None: + results.append(classify_image(img, model, cfg)) + continue + family, scores = run_inference(img, model, cfg) + tags = postprocess_scores(family, scores, model, cfg) + results.append(tags) + try: + entry = _debug_entry(model, cfg) + entry["per_frame"].append( + _debug_pairs(candidate_scores(family, scores, model, cfg)) + ) + _finish_debug_entry(entry, tags) + debug.append(entry) + except Exception as exc: + stashlog.log.error( + f"Debug candidates for {model.get('name')} failed: {exc}" + ) + except Exception as exc: + stashlog.log.error(f"Model {model.get('name')} failed: {exc}") + if not results: + return [] + if len(results) == 1: + return results[0] + return merge_tag_results(results) + + +def run_model(path, models, cfg, debug=None): + """Classify a single image path (supports .zip members).""" + return classify_with_models(load_pil_image(path), models, cfg, debug=debug) + + +def run_model_batch(paths, models, cfg, debug_out=None): + """Classify a list of image paths, returning one result list per path. + + When ``debug_out`` is a list, one debug record per path is appended to it. + """ + results = [] + for path in paths: + debug = [] if debug_out is not None else None + try: + results.append(run_model(path, models, cfg, debug=debug)) + except Exception as exc: + stashlog.log.error(f"Failed to process image {path}: {exc}") + results.append([]) + if debug_out is not None: + debug_out.append(debug) + return results + + +def _record_hits(hits_out, model, per_frame): + """Store one model's accepted per-frame tags for the marker builder.""" + if hits_out is not None: + hits_out[model.get("name")] = per_frame + + +def _infer_frames_one_model(images, model, cfg, debug=None, hits_out=None): + """Aggregate one model's per-frame results into a scene-level tag list. + + When ``debug`` is a dict the per-frame candidate confidences are stored + under ``debug[model name]`` for the calibration report. When ``hits_out`` + is a dict the accepted per-frame tags are stored under the model name, one + row per frame -- a failed frame yields an empty row, so the rows always + stay aligned with ``images`` (the video pipeline uses them to place scene + markers). + """ + entry = None + if debug is not None: + entry = _debug_entry(model, cfg) + debug[model.get("name")] = entry + + if model["family"] == "yolo": + try: + batch_size = max( + 1, int(cfg.get("infer_batch_size", INFER_BATCH_SIZE_DEFAULT)) + ) + except (TypeError, ValueError): + batch_size = INFER_BATCH_SIZE_DEFAULT + + per_frame = [] + for start in range(0, len(images), batch_size): + chunk = images[start:start + batch_size] + expected = len(per_frame) + len(chunk) + try: + batch = yolo_preprocess_batch(chunk, model["size"]) + outputs = model["session"].run( + [model["output_name"]], {model["input_name"]: batch} + )[0] + except Exception as exc: + stashlog.log.error(f"Failed to classify {len(chunk)} frame(s): {exc}") + else: + for row in outputs: + probs = _as_probabilities(np.asarray(row).reshape(-1)) + per_frame.append( + yolo_postprocess( + probs, + model, + thresholds=model.get("thresholds"), + margin=cfg["yolo_margin"], + default_threshold=cfg["yolo_threshold"], + top_k=cfg.get("yolo_top_k", YOLO_TOP_K_DEFAULT), + ) + ) + if entry is not None: + entry["per_frame"].append( + _debug_pairs(candidate_scores("yolo", probs, model, cfg)) + ) + # A failed batch must not shift the frame indices. + while len(per_frame) < expected: + per_frame.append([]) + if entry is not None: + entry["per_frame"].append([]) + + tags = yolo_aggregate( + per_frame, + frame_count=len(images), + multi=cfg.get("yolo_multi", True), + min_votes=cfg.get("yolo_min_votes", YOLO_MIN_VOTES_DEFAULT), + share=cfg.get("yolo_share", YOLO_SHARE_DEFAULT), + ) + _finish_debug_entry(entry, tags) + _record_hits(hits_out, model, per_frame) + return tags + + per_frame = [] + for img in images: + try: + if entry is None: + per_frame.append(classify_image(img, model, cfg)) + continue + family, scores = run_inference(img, model, cfg) + tags = postprocess_scores(family, scores, model, cfg) + per_frame.append(tags) + entry["per_frame"].append( + _debug_pairs(candidate_scores(family, scores, model, cfg)) + ) + except Exception as exc: + stashlog.log.error(f"Failed to classify frame: {exc}") + per_frame.append([]) + if entry is not None: + entry["per_frame"].append([]) + + if model["family"] in ("yolo-det", "sexpos"): + # These models name classes per accepted frame, so the votes are + # pooled exactly like the single-label classifier's accepted frames. + tags = yolo_aggregate( + per_frame, + frame_count=len(images), + multi=cfg.get("yolo_multi", True), + min_votes=cfg.get("yolo_min_votes", YOLO_MIN_VOTES_DEFAULT), + share=cfg.get("yolo_share", YOLO_SHARE_DEFAULT), + ) + else: + tags = wd_aggregate(per_frame) + _finish_debug_entry(entry, tags) + _record_hits(hits_out, model, per_frame) + return tags + + +def infer_frames(frame_paths, models, cfg, debug=None, hits_out=None): + """Classify every frame with every model and merge the tags. + + Frames are read once. YOLO classification models are pushed through the + session in batches (their batch dimension is dynamic), which keeps the GPU + busy instead of paying per-frame call overhead; the other families (WD and + JoyTag have a fixed batch of 1, the detector and sexpos models run one + frame at a time) are handled per image. With several models the per-model + scene tags are merged, keeping the best score for each tag. + + When ``debug`` is a dict the per-frame candidate confidences of every model + are collected in it (used by the calibration report); when ``hits_out`` is + a dict the accepted per-frame tags of every model are collected in it + (used to place scene markers). + """ + if isinstance(models, dict): + models = [models] + + images = [] + for path in frame_paths: + try: + images.append(Image.open(path).convert("RGB")) + except Exception as exc: + stashlog.log.error(f"Failed to load frame {path}: {exc}") + + results = [] + for model in models: + try: + results.append( + _infer_frames_one_model( + images, model, cfg, debug=debug, hits_out=hits_out + ) + ) + except Exception as exc: + stashlog.log.error(f"Model {model.get('name')} failed on frames: {exc}") + + if not results: + return [] + if len(results) == 1: + return results[0] + return merge_tag_results(results) + + +def run_model_frames(frame_paths, models, cfg, debug=None, hits_out=None): + """Alias kept for readability in the video pipeline.""" + return infer_frames(frame_paths, models, cfg, debug=debug, hits_out=hits_out) diff --git a/plugins/tagAI/tagai_core/jobs.py b/plugins/tagAI/tagai_core/jobs.py new file mode 100755 index 00000000..5f5ab698 --- /dev/null +++ b/plugins/tagAI/tagai_core/jobs.py @@ -0,0 +1,102 @@ +"""Pagination loop and image/scene job runners.""" + +from tagai_core import stash_client, stashlog +from tagai_core.debug import debug_add_item +from tagai_core.images import apply_tags +from tagai_core.inference import run_model_batch +from tagai_core.scenes import process_scene +from tagai_core.tags import filter_tags +import traceback + + +def iterate_job(fetch_page, process_batch, mode, per_page, on_progress=None): + """Pagination loop shared by images and scenes. + + ``fetch_page(page) -> (items, count)`` and ``process_batch(items)``. + + * ``TagAll`` / ``Tagged``: the result set *shrinks* as items get tagged, so + always fetch page 1; stop when the page is empty or no returned item is + new (this is the fix for the historic infinite re-tag loop). + * ``ForceTag``: the result set is stable, so advance pages until a page is + empty (with a safety bound based on ``count``). + + Returns the set of processed ids. + """ + processed = set() + page = 1 + total = 0 + + while True: + items, count = fetch_page(page) + if not items: + break + + fresh = [it for it in items if it.get("id") not in processed] + + if mode in ("TagAll", "Tagged"): + if not fresh: + break + if fresh: + process_batch(fresh) + processed.update(it["id"] for it in fresh) + total += len(fresh) + if on_progress and count: + on_progress(min(1.0, total / float(count))) + + if mode == "ForceTag": + page += 1 + if count and page * per_page > count + per_page: + break + else: + page = 1 + + return processed + + +def run_image_job(fetch_page, mode, per_page, models, cfg, + flag_tag_id, ai_tagged_id, prefix): + def process_batch(items): + valid = [it for it in items if it.get("files")] + paths = [it["files"][0]["path"] for it in valid] + debug_out = [] if cfg.get("debug") else None + tag_batches = run_model_batch(paths, models, cfg, debug_out=debug_out) + for index, (it, tags) in enumerate(zip(valid, tag_batches)): + try: + tags = filter_tags( + tags, cfg.get("tag_whitelist"), cfg.get("tag_blacklist") + ) + ids = [tid for tid in (stash_client.ensure_tag(prefix + name) + for name, _score in tags) if tid] + apply_tags(it["id"], ids, flag_tag_id, ai_tagged_id) + if debug_out is not None: + debug_add_item({ + "id": it.get("id"), + "path": paths[index], + "duration": None, + "frames": 1, + "models": debug_out[index], + "applied_tags": [prefix + name for name, _score in tags], + }) + except Exception as exc: + stashlog.log.error(f"Image {it.get('id')} failed: {exc}\n{traceback.format_exc()}") + + return iterate_job(fetch_page, process_batch, mode, per_page, + on_progress=lambda p: stashlog.log.progress(p)) + + +def run_scene_job(fetch_page, mode, per_page, models, cfg, + flag_tag_id, ai_tagged_id, prefix): + def process_batch(items): + for it in items: + try: + process_scene(it, models, cfg, flag_tag_id, ai_tagged_id, prefix) + except Exception as exc: + stashlog.log.error(f"Scene {it.get('id')} failed: {exc}\n{traceback.format_exc()}") + + return iterate_job(fetch_page, process_batch, mode, per_page, + on_progress=lambda p: stashlog.log.progress(p)) + + +# ========================================================================== # +# 7. ARG PARSING & DISPATCHER +# ========================================================================== # diff --git a/plugins/tagAI/tagai_core/markers.py b/plugins/tagAI/tagai_core/markers.py new file mode 100755 index 00000000..eead0293 --- /dev/null +++ b/plugins/tagAI/tagai_core/markers.py @@ -0,0 +1,263 @@ +"""Scene marker planning and creation.""" + +from tagai_core import stash_client, stashlog +from tagai_core.tags import filter_marker_tags +from tagai_core.config import MARKER_MAX_DEFAULT, MARKER_MIN_LENGTH + + +def get_scene_markers(scene_id): + """Existing markers of a scene as ``[{primary_tag_id, seconds}]``.""" + q = """query($id:ID!){ + findScene(id:$id){ scene_markers{ seconds primary_tag{ id } } } + }""" + result = stash_client.gql(q, {"id": scene_id}) + markers = ( + ((result.get("data") or {}).get("findScene") or {}).get("scene_markers") + or [] + ) + existing = [] + for marker in markers: + primary = marker.get("primary_tag") or {} + existing.append({ + "primary_tag_id": primary.get("id"), + "seconds": float(marker.get("seconds") or 0.0), + }) + return existing + + +def create_scene_marker(scene_id, primary_tag_id, start, end, title, + tag_ids=None): + """Create one scene marker; returns whether Stash accepted it.""" + q = """mutation($input:SceneMarkerCreateInput!){ + sceneMarkerCreate(input:$input){ id } + }""" + payload = { + "scene_id": scene_id, + "primary_tag_id": primary_tag_id, + "seconds": round(float(start), 3), + "end_seconds": round(float(end), 3), + "title": title, + } + if tag_ids: + payload["tag_ids"] = list(tag_ids) + result = stash_client.gql(q, {"input": payload}) + created = ((result.get("data") or {}).get("sceneMarkerCreate") or {}) + return bool(created.get("id")) + + +def parse_marker_mode(value): + """Normalise the ``markers`` argument -> ``off``, ``ranges`` or ``best``.""" + text = str(value or "").strip().lower() + if text in ("", "off", "false", "0", "no", "none", "disabled"): + return "off" + if text in ("ranges", "range", "runs"): + return "ranges" + if text in ("best", "top", "peak"): + return "best" + stashlog.log.error(f"Unknown markers mode {value!r}; expected off, ranges or best") + return "off" + + +def _frame_spacing(times): + """Median gap between consecutive sampled frames (0.0 when unknown).""" + if not times or len(times) < 2: + return 0.0 + gaps = sorted(float(b) - float(a) for a, b in zip(times, times[1:])) + return gaps[len(gaps) // 2] + + +def combine_hits(hits_out, frame_count): + """Merge per-model per-frame hits -> one ``{tag: score}`` map per frame.""" + combined = [{} for _ in range(max(0, int(frame_count)))] + for per_frame in (hits_out or {}).values(): + for index, row in enumerate(per_frame or []): + if index >= len(combined): + break + for name, score in row or []: + best = combined[index].get(name) + if best is None or score > best: + combined[index][name] = float(score) + return combined + + +def _merge_run_rows(run, hits): + """Best score per tag over all frames of one run.""" + merged = {} + for index, _moment, _score in run: + for tag, score in (hits[index] or {}).items(): + if tag not in merged or score > merged[tag]: + merged[tag] = float(score) + return merged + + +def _marker_span(name, start, end, score, wanted, row, pad, duration): + """One marker record, padded and clamped to the video's duration.""" + start = max(0.0, float(start) - pad) + end = float(end) + pad + end = max(end, start + MARKER_MIN_LENGTH) + if duration: + end = min(float(duration), end) + start = min(start, max(0.0, end - MARKER_MIN_LENGTH)) + return { + "tag": name, + "start": round(start, 3), + "end": round(end, 3), + "score": round(float(score), 4), + "tags": [tag for tag in wanted if tag != name and tag in row], + } + + +def build_markers(tags, hits, frame_times, duration, mode="ranges", pad=0.0, + gap=0.0): + """Turn the per-frame hits of a scene into marker spans. + + ``hits`` holds one ``{tag: score}`` mapping per sampled frame and + ``frame_times`` the matching timestamps (both in sampling order, so index + *i* of one belongs to index *i* of the other). Only the scene-level + ``tags`` are considered, so the markers always agree with the applied tags. + + With ``mode="ranges"`` consecutive hits of a tag become a single marker + spanning from its first to its last hit; with ``mode="best"`` only the most + confident frame of each tag is used. Every span is padded by ``pad`` + seconds on both sides -- so a marker lasts at least ``2 * pad`` seconds -- + and clamped to ``[0, duration]``. The other tags that fired during the same + span are returned as co-tags. + + Two hits of one tag are also merged when they are no further apart than + ``2 * pad + gap`` seconds, so a tag that fires in bursts produces one + marker per burst instead of one per sample. ``gap`` defaults to ``0``, + which only merges hits whose padded spans would already touch. + """ + wanted = [name for name, _score in tags] + if not wanted or not hits or not frame_times: + return [] + + try: + pad = max(0.0, float(pad)) + except (TypeError, ValueError): + pad = 0.0 + + try: + gap = max(0.0, float(gap)) + except (TypeError, ValueError): + gap = 0.0 + + points = {} + for index in range(min(len(hits), len(frame_times))): + row = hits[index] or {} + for name in wanted: + if name in row: + points.setdefault(name, []).append( + (index, float(frame_times[index]), float(row[name])) + ) + + markers = [] + for name in wanted: + found = points.get(name) or [] + if not found: + continue + + if mode == "best": + index, moment, score = max(found, key=lambda item: (item[2], -item[0])) + markers.append( + _marker_span(name, moment, moment, score, wanted, + hits[index], pad, duration) + ) + continue + + run = [found[0]] + spans = [] + for point in found[1:]: + # Adjacent samples belong to the same moment; so do hits whose + # padded spans would overlap (or come within the merge gap). + if (point[0] - run[-1][0] == 1 + or point[1] - run[-1][1] <= 2 * pad + gap + 1e-6): + run.append(point) + else: + spans.append(run) + run = [point] + spans.append(run) + + for span in spans: + markers.append( + _marker_span(name, span[0][1], span[-1][1], + max(point[2] for point in span), wanted, + _merge_run_rows(span, hits), pad, duration) + ) + + markers.sort(key=lambda marker: (marker["start"], marker["tag"])) + return markers + + +def create_scene_markers(scene, tags, hits, frame_times, duration, cfg, prefix): + """Create the scene markers for one scene; returns the created spans. + + The marker whitelist/blacklist is applied first, so a filtered tag never + becomes a marker (or the co-tag of one). Re-runs do not stack markers: a + span is skipped when the scene already has a marker with the same primary + tag close to (within one padding) the same moment. At most ``marker_max`` + markers are created per scene. + """ + mode = cfg.get("markers") or "off" + if mode == "off": + return [] + + tags = filter_marker_tags( + tags, cfg.get("marker_whitelist"), cfg.get("marker_blacklist") + ) + + pad = cfg.get("marker_pad") + if pad is None: + pad = _frame_spacing(frame_times) + try: + gap = max(0.0, float(cfg.get("marker_gap") or 0.0)) + except (TypeError, ValueError): + gap = 0.0 + + spans = build_markers(tags, hits, frame_times, duration, mode=mode, pad=pad, + gap=gap) + if not spans: + return [] + + try: + limit = max(0, int(cfg.get("marker_max", MARKER_MAX_DEFAULT))) + except (TypeError, ValueError): + limit = MARKER_MAX_DEFAULT + if limit == 0: + return [] + + scene_id = scene.get("id") + existing = get_scene_markers(scene_id) + tolerance = max(MARKER_MIN_LENGTH, 2.0 * float(pad)) + + created = [] + for span in spans: + if len(created) >= limit: + stashlog.log.info( + f"Scene {scene_id}: marker limit ({limit}) reached, " + "not creating further markers" + ) + break + + tag_id = stash_client.ensure_tag(prefix + span["tag"]) + if not tag_id: + continue + if any( + marker.get("primary_tag_id") == tag_id + and abs(marker.get("seconds", 0.0) - span["start"]) <= tolerance + for marker in existing + ): + continue + + co_tag_ids = [ + tid for tid in (stash_client.ensure_tag(prefix + name) for name in span["tags"]) + if tid and tid != tag_id + ] + if create_scene_marker(scene_id, tag_id, span["start"], span["end"], + span["tag"], co_tag_ids): + existing.append({"primary_tag_id": tag_id, "seconds": span["start"]}) + created.append(span) + + if created: + stashlog.log.info(f"Scene {scene_id}: created {len(created)} marker(s)") + return created diff --git a/plugins/tagAI/tagai_core/models.py b/plugins/tagAI/tagai_core/models.py new file mode 100755 index 00000000..2ee0d11b --- /dev/null +++ b/plugins/tagAI/tagai_core/models.py @@ -0,0 +1,423 @@ +"""Model registry: ONNX session cache, Hugging Face download, labels.""" + +from tagai_core import config, stashlog +from tagai_core.config import ( + HF_ENDPOINT, HF_TOKEN_ENV_VARS, HF_TOKEN_FILENAMES, JOYTAG_SIZE, MODELS, + SEXPOS_CLASS_NAMES, SEXPOS_SIZE, YOLO_CLASS_NAMES, YOLO_DEFAULT_THRESHOLDS, + YOLO_DET_CLASS_NAMES, YOLO_DET_SIZE, YOLO_ZH_TO_EN, +) +from tagai_core.config import requests +import ast +import csv +import json +import os + + +_SESSION_CACHE = {} # model name -> entry dict + + +def _resolve_wd_size(shape): + """WD taggers are NHWC [N, H, W, 3]; return H (fall back to 448).""" + if len(shape) == 4 and isinstance(shape[2], int): + return shape[2] + if len(shape) == 4 and isinstance(shape[1], int): + return shape[1] + return 448 + + +def load_tags(csv_path): + """Parse a WD ``selected_tags.csv``. + + Returns ``(tags, general_idx, char_idx)`` where the rows are laid out as + ``[general..., ratings..., characters...]``. ``char_idx`` is ``None`` when + the file has no character rows. + """ + tags = [] + general_idx = None + char_idx = None + + with open(csv_path, encoding="utf-8") as f: + reader = csv.reader(f) + next(reader, None) # header + for row in reader: + if len(row) < 3: + continue + name = row[1].replace("_", " ") + tags.append(name) + + if general_idx is None and row[2] == "0": + general_idx = len(tags) - 1 + elif char_idx is None and row[2] == "4": + char_idx = len(tags) - 1 + + if general_idx is None: + raise RuntimeError( + f"CSV boundary detection failed for {csv_path}: no rows with category '0'" + ) + + return tags, general_idx, char_idx + + +def load_joytag_tags(txt_path): + """Read JoyTag's ``top_tags.txt`` (one tag per line, index-ordered).""" + with open(txt_path, encoding="utf-8") as f: + return [line.strip() for line in f if line.strip()] + + +def load_yolo_class_names(session): + """Read class names from ONNX metadata (``names``), else the built-in list. + + The upstream model ships Chinese names; they are translated to English so + the resulting Stash tags stay readable. + """ + try: + meta = session.get_modelmeta().custom_metadata_map + raw = meta.get("names") + if raw: + parsed = ast.literal_eval(raw) + if isinstance(parsed, dict): + names = [str(parsed[k]) for k in sorted(parsed, key=lambda x: int(x))] + elif isinstance(parsed, (list, tuple)): + names = [str(x) for x in parsed] + else: + names = None + if names: + return [YOLO_ZH_TO_EN.get(n, n) for n in names] + except Exception as exc: # pragma: no cover - depends on model metadata + stashlog.log.error(f"Could not read YOLO class names from metadata: {exc}") + return list(YOLO_CLASS_NAMES) + + +def load_det_class_names(session): + """Read detection class names from ONNX metadata (``names``), else built-in. + + The YOLO11-NSFW export stores ``names`` as a ``{index: name}`` mapping in + its metadata; when it is missing the bundled list is used. + """ + try: + meta = session.get_modelmeta().custom_metadata_map + raw = meta.get("names") + if raw: + parsed = ast.literal_eval(raw) + if isinstance(parsed, dict): + return [str(parsed[k]) for k in sorted(parsed, key=lambda x: int(x))] + if isinstance(parsed, (list, tuple)): + return [str(x) for x in parsed] + except Exception as exc: # pragma: no cover - depends on model metadata + stashlog.log.error(f"Could not read detection class names from metadata: {exc}") + return list(YOLO_DET_CLASS_NAMES) + + +def load_yolo_thresholds(model_name): + """Return per-class thresholds for a YOLO model, else ``None``. + + Starts from the thresholds published by the model author (bundled in + :data:`YOLO_DEFAULT_THRESHOLDS`) and lets an optional + ``Models/-thresholds.json`` override them. Keys in that file + may be English or the model's original Chinese names. + """ + spec = MODELS.get(model_name, {}) + if spec.get("family") != "yolo": + return None + + thresholds = dict(YOLO_DEFAULT_THRESHOLDS) + + onnx_name = spec.get("onnx") + if onnx_name: + base = os.path.splitext(onnx_name)[0] + path = os.path.join(config.MODEL_DIR, base + "-thresholds.json") + if os.path.isfile(path): + try: + with open(path, encoding="utf-8") as f: + data = json.load(f) + except (OSError, ValueError) as exc: + stashlog.log.error(f"Failed to read YOLO thresholds {path}: {exc}") + data = None + + if isinstance(data, dict): + for k, v in data.items(): + try: + key = YOLO_ZH_TO_EN.get(str(k), str(k)) + thresholds[key] = float(v) + except (TypeError, ValueError): + continue + + return thresholds or None + + +def resolve_hf_token(explicit=None): + """Return a Hugging Face token from the arg, env vars, or a local file.""" + if explicit: + return str(explicit).strip() + + for var in HF_TOKEN_ENV_VARS: + value = os.environ.get(var) + if value and value.strip(): + return value.strip() + + for name in HF_TOKEN_FILENAMES: + path = os.path.join(config.PLUGIN_DIR, name) + if os.path.isfile(path): + try: + with open(path, encoding="utf-8") as f: + token = f.readline().strip() + if token: + return token + except OSError: + pass + return None + + +def _hf_url(repo, path_in_repo): + return f"{HF_ENDPOINT}/{repo}/resolve/main/{path_in_repo}" + + +def download_file(url, dest, token=None, on_progress=None): + """Stream ``url`` to ``dest`` atomically (via a ``.part`` temp file).""" + headers = {"User-Agent": "tagAI"} + if token: + headers["Authorization"] = f"Bearer {token}" + + os.makedirs(os.path.dirname(dest) or ".", exist_ok=True) + tmp = dest + ".part" + + with requests.get( + url, headers=headers, stream=True, timeout=60, allow_redirects=True + ) as resp: + if resp.status_code in (401, 403): + raise PermissionError( + f"access denied ({resp.status_code}) for {url}" + ) + resp.raise_for_status() + + try: + total = int(resp.headers.get("Content-Length") or 0) + except (TypeError, ValueError): + total = 0 + + done = 0 + with open(tmp, "wb") as fh: + for chunk in resp.iter_content(chunk_size=1024 * 1024): + if not chunk: + continue + fh.write(chunk) + done += len(chunk) + if on_progress and total: + on_progress(min(1.0, done / total)) + + os.replace(tmp, dest) + return dest + + +def ensure_model_files(name, token=None, allow_download=True, on_progress=None): + """Make sure every file for ``name`` exists locally, downloading if needed. + + Returns the list of files that were downloaded (empty when everything was + already present). Raises a clear error when a file is missing and downloads + are disabled or denied. + """ + if name not in MODELS: + raise ValueError( + f"Unknown model {name!r}. Available: {', '.join(sorted(MODELS))}" + ) + + spec = MODELS[name] + files = spec.get("files") or [] + missing = [ + (local, remote) + for local, remote in files + if not os.path.isfile(os.path.join(config.MODEL_DIR, local)) + ] + optional = [ + (local, remote) + for local, remote in (spec.get("optional_files") or []) + if not os.path.isfile(os.path.join(config.MODEL_DIR, local)) + ] + if not missing and not optional: + return [] + + if not allow_download: + if missing: + raise FileNotFoundError( + "Missing model file(s): " + + ", ".join(local for local, _ in missing) + + "\nSee Models/README.md for download instructions." + ) + # Only optional files are missing; skip them (bundled defaults apply). + return [] + + repo = spec.get("repo") + if repo and (missing or optional): + token = resolve_hf_token(token) + else: + token = None + + if missing and not repo: + raise FileNotFoundError( + "No download source configured for model " + f"{name!r}; place the files in {config.MODEL_DIR} manually " + "(see Models/README.md)." + ) + + if repo and spec.get("gated") and (missing or optional) and not token: + stashlog.log.warning( + f"Model {name!r} is gated on Hugging Face. Downloading without a " + "token will likely fail; set HF_TOKEN (or a local hf_token file) " + "after accepting the model's terms. See Models/README.md." + ) + + downloaded = [] + + def _fetch(local, remote): + dest = os.path.join(config.MODEL_DIR, local) + url = _hf_url(repo, remote) + stashlog.log.info(f"Downloading {local} from {url}") + + def _report(fraction, _name=local): + if on_progress: + try: + on_progress(fraction) + except Exception: + pass + stashlog.log.progress(fraction) + + download_file(url, dest, token=token, on_progress=_report) + return dest + + for local, remote in missing: + try: + dest = _fetch(local, remote) + except PermissionError as exc: + raise FileNotFoundError( + f"Could not download {local} from {remote}: {exc}. If this is a " + "gated model, accept its terms on Hugging Face and provide a " + "token via HF_TOKEN or a local hf_token file. " + "See Models/README.md." + ) from exc + downloaded.append(dest) + stashlog.log.info(f"Saved {dest}") + + # Optional files (e.g. the published threshold table) are best-effort. + for local, remote in optional: + if not repo: + break + try: + dest = _fetch(local, remote) + except Exception as exc: + stashlog.log.warning(f"Could not download optional file {local}: {exc}") + continue + downloaded.append(dest) + stashlog.log.info(f"Saved {dest}") + + return downloaded + + +def get_model(name): + """Return a cached model entry, loading it on first use. + + Raises clear errors when the model name is unknown or a required file is + missing so the plugin fails gracefully (Stash must not crash). + """ + if name in _SESSION_CACHE: + return _SESSION_CACHE[name] + + if name not in MODELS: + raise ValueError( + f"Unknown model {name!r}. Available: {', '.join(sorted(MODELS))}" + ) + + spec = MODELS[name] + onnx_path = os.path.join(config.MODEL_DIR, spec["onnx"]) + if not os.path.isfile(onnx_path): + raise FileNotFoundError( + f"Model file not found: {onnx_path}\n" + "See Models/README.md for download instructions." + ) + + try: + import onnxruntime as ort + except ImportError as exc: # pragma: no cover - depends on env + raise ImportError( + "onnxruntime is required for TagAI. Install with: " + "pip install -r requirements.txt" + ) from exc + + session = ort.InferenceSession( + onnx_path, + providers=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) + input_meta = session.get_inputs()[0] + output_name = session.get_outputs()[0].name + + entry = { + "name": name, + "session": session, + "family": spec["family"], + "input_name": input_meta.name, + "output_name": output_name, + "spec": spec, + } + + if spec["family"] == "wd": + csv_path = os.path.join(config.MODEL_DIR, spec["labels"]) + if not os.path.isfile(csv_path): + raise FileNotFoundError( + f"Label CSV not found: {csv_path}\n" + "See Models/README.md for download instructions." + ) + tags, general_idx, char_idx = load_tags(csv_path) + entry.update( + { + "tags": tags, + "general_idx": general_idx, + "char_idx": char_idx, + "size": _resolve_wd_size(input_meta.shape), + } + ) + elif spec["family"] == "yolo": + entry.update( + { + "size": 224, + "class_names": load_yolo_class_names(session), + "thresholds": load_yolo_thresholds(name), + } + ) + elif spec["family"] == "yolo-det": + entry.update( + { + "size": YOLO_DET_SIZE, + "class_names": load_det_class_names(session), + } + ) + elif spec["family"] == "sexpos": + entry.update( + { + "size": SEXPOS_SIZE, + "class_names": list(SEXPOS_CLASS_NAMES), + } + ) + + else: # joytag + csv_path = os.path.join(config.MODEL_DIR, spec["labels"]) + if not os.path.isfile(csv_path): + raise FileNotFoundError( + f"Tag list not found: {csv_path}\n" + "See Models/README.md for download instructions." + ) + entry.update( + { + "tags": load_joytag_tags(csv_path), + "size": JOYTAG_SIZE, + } + ) + + _SESSION_CACHE[name] = entry + return entry + + +def unload_model(name): + """Drop a cached session (frees GPU memory).""" + _SESSION_CACHE.pop(name, None) + + +# ------------------------------ preprocessing ----------------------------- # diff --git a/plugins/tagAI/tagai_core/postprocess.py b/plugins/tagAI/tagai_core/postprocess.py new file mode 100755 index 00000000..c1807914 --- /dev/null +++ b/plugins/tagAI/tagai_core/postprocess.py @@ -0,0 +1,348 @@ +"""Per-family score post-processing and frame aggregation.""" + +from tagai_core.config import ( + CHAR_THRESHOLD, JOYTAG_THRESHOLD_DEFAULT, SEXPOS_CLASS_NAMES, SEXPOS_MARGIN_DEFAULT, + SEXPOS_THRESHOLD_DEFAULT, TAG_THRESHOLD, YOLO_DET_CLASS_NAMES, + YOLO_DET_CONF_DEFAULT, YOLO_DET_IOU_DEFAULT, YOLO_MARGIN_DEFAULT, + YOLO_MIN_VOTES_DEFAULT, YOLO_SHARE_DEFAULT, YOLO_THRESHOLD_DEFAULT, + YOLO_TOP_K_DEFAULT, +) +from tagai_core.config import np +import math + + +def softmax(x): + x = np.asarray(x, dtype=np.float32) + x = x - np.max(x) + e = np.exp(x) + s = float(np.sum(e)) + if s == 0.0: + return e + return e / s + + +def _as_probabilities(out): + """Return a probability vector. + + Some classification ONNX exports emit raw logits and must be softmaxed by + the caller, while others (e.g. yolo26n-cls-porn, which contains a Softmax + node) already emit probabilities. If the output already looks softmaxed + (all in [0, 1] and summing to ~1) it is returned unchanged. + """ + arr = np.asarray(out, dtype=np.float32).reshape(-1) + if arr.size: + if ( + float(arr.min()) >= 0.0 + and float(arr.max()) <= 1.0 + and abs(float(arr.sum()) - 1.0) < 1e-3 + ): + return arr + return softmax(arr) + + +def _as_sigmoid_probabilities(out): + """Return per-tag probabilities for a sigmoid (multi-label) head. + + JoyTag emits raw logits that must be passed through a sigmoid, but some + exports already include the activation. If every value is already inside + [0, 1] the output is returned unchanged. + """ + arr = np.asarray(out, dtype=np.float32).reshape(-1) + if arr.size and float(arr.min()) >= 0.0 and float(arr.max()) <= 1.0: + return arr + return 1.0 / (1.0 + np.exp(-arr)) + + +def _nms(boxes, scores, iou_threshold): + """Greedy non-maximum suppression -> indices of the kept boxes.""" + boxes = np.asarray(boxes, dtype=np.float32) + scores = np.asarray(scores, dtype=np.float32) + if boxes.size == 0: + return [] + + x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] + areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1) + order = scores.argsort()[::-1] + + keep = [] + while order.size: + i = int(order[0]) + keep.append(i) + if order.size == 1: + break + rest = order[1:] + xx1 = np.maximum(x1[i], x1[rest]) + yy1 = np.maximum(y1[i], y1[rest]) + xx2 = np.minimum(x2[i], x2[rest]) + yy2 = np.minimum(y2[i], y2[rest]) + inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) + union = areas[i] + areas[rest] - inter + iou = np.divide(inter, union, out=np.zeros_like(inter), where=union > 0) + order = rest[iou <= iou_threshold] + + return keep + + +def yolo_det_postprocess(output, model, + conf_threshold=YOLO_DET_CONF_DEFAULT, + iou_threshold=YOLO_DET_IOU_DEFAULT): + """Decode a YOLO11 detection head -> ``[(class_name, confidence)]``. + + The exported graph ends at the raw head, so boxes are decoded here: the + output is ``[4 + num_classes, anchors]`` (or the transposed layout), each + anchor carrying a centre-based ``xywh`` box followed by one score per + class. Only class names and confidences are returned -- box positions are + not needed for tagging -- and overlapping boxes are removed with NMS. + """ + preds = np.asarray(output, dtype=np.float32) + while preds.ndim > 2: + preds = preds[0] + if preds.ndim != 2: + return [] + + names = model.get("class_names") or list(YOLO_DET_CLASS_NAMES) + num_classes = len(names) + + # Rows must be anchors: accept either [anchors, 4 + nc] or [4 + nc, anchors]. + if preds.shape[1] != 4 + num_classes and preds.shape[0] == 4 + num_classes: + preds = preds.T + if preds.shape[1] < 5: + return [] + num_classes = min(num_classes, preds.shape[1] - 4) + if num_classes <= 0: + return [] + + boxes_xywh = preds[:, :4] + scores = preds[:, 4:4 + num_classes] + class_ids = scores.argmax(axis=1) + confs = scores[np.arange(scores.shape[0]), class_ids] + + mask = confs >= conf_threshold + if not np.any(mask): + return [] + boxes_xywh = boxes_xywh[mask] + confs = confs[mask] + class_ids = class_ids[mask] + + cx, cy, bw, bh = (boxes_xywh[:, i] for i in range(4)) + boxes_xyxy = np.stack( + [cx - bw / 2.0, cy - bh / 2.0, cx + bw / 2.0, cy + bh / 2.0], axis=1 + ) + + keep = _nms(boxes_xyxy, confs, iou_threshold) + + results = [] + for i in keep: + cid = int(class_ids[i]) + name = names[cid] if cid < len(names) else f"class_{cid}" + results.append((name, float(confs[i]))) + return results + + +def _sexpos_from_probs(probs, model, threshold=SEXPOS_THRESHOLD_DEFAULT, + margin=SEXPOS_MARGIN_DEFAULT): + """Accept at most the top class of one frame -> ``[(name, score)]``. + + A frame is only tagged when the top class reaches ``threshold`` and, when + ``margin`` is non-zero, beats the runner-up by at least that much. + """ + probs = np.asarray(probs, dtype=np.float32).reshape(-1) + if probs.size == 0: + return [] + + names = model.get("class_names") or list(SEXPOS_CLASS_NAMES) + if len(names) != probs.size: + names = [f"class_{i}" for i in range(probs.size)] + + order = np.argsort(probs)[::-1] + top1 = float(probs[order[0]]) + top2 = float(probs[order[1]]) if order.size > 1 else 0.0 + if top1 < threshold: + return [] + if margin and (top1 - top2) < margin: + return [] + return [(names[int(order[0])], top1)] + + +def sexpos_postprocess(logits, model, threshold=SEXPOS_THRESHOLD_DEFAULT, + margin=SEXPOS_MARGIN_DEFAULT): + """Accept at most the top class of one frame -> ``[(name, score)]``. + + The exported graph ends at the raw logits, so a softmax is applied here. + A frame is only tagged when the top class reaches ``threshold`` and, when + ``margin`` is non-zero, beats the runner-up by at least that much. + """ + return _sexpos_from_probs( + softmax(np.asarray(logits, dtype=np.float32).reshape(-1)), + model, + threshold=threshold, + margin=margin, + ) + + +def wd_postprocess(probs, model, tag_threshold=TAG_THRESHOLD, + char_threshold=CHAR_THRESHOLD, tag_characters=False): + """Return ``[(name, score)]`` for the WD model given a probability vector.""" + tags = model["tags"] + general_idx = model["general_idx"] + char_idx = model["char_idx"] + if char_idx is None: + char_idx = len(tags) + + probs = np.asarray(probs).reshape(-1) + + general = [ + (tags[i], float(probs[i])) + for i in range(general_idx, min(char_idx, len(tags))) + if probs[i] > tag_threshold + ] + chars = [] + if tag_characters: + chars = [ + (tags[i], float(probs[i])) + for i in range(char_idx, len(tags)) + if probs[i] > char_threshold + ] + + return general + chars + + +def joytag_postprocess(probs, model, threshold=JOYTAG_THRESHOLD_DEFAULT): + """Return ``[(name, score)]`` for JoyTag: every tag above ``threshold``.""" + tags = model.get("tags") or [] + probs = np.asarray(probs, dtype=np.float32).reshape(-1) + if len(tags) != len(probs): + tags = [f"tag_{i}" for i in range(len(probs))] + + hits = [ + (tags[i], float(probs[i])) + for i in range(len(probs)) + if probs[i] > threshold + ] + hits.sort(key=lambda kv: kv[1], reverse=True) + return hits + + +def yolo_postprocess(probs, model, thresholds=None, + margin=YOLO_MARGIN_DEFAULT, + default_threshold=YOLO_THRESHOLD_DEFAULT, + top_k=YOLO_TOP_K_DEFAULT): + """Accept up to ``top_k`` classes for one frame -> ``[(name, score)]``. + + The frame is only trusted when the top class passes the author's rule: + + ``conf_ok = top1 >= threshold`` (per-class if available) + ``margin_ok = top1 - top2 >= margin`` + + With ``top_k > 1`` the runner-up classes are additionally returned for a + trusted frame when they clear their own confidence threshold. This raises + recall on frames that show more than one act, at the cost of potential + false positives; the default ``top_k = 1`` keeps the original behaviour. + """ + probs = np.asarray(probs, dtype=np.float32).reshape(-1) + names = model.get("class_names") or [] + if len(names) != len(probs): + names = [f"class_{i}" for i in range(len(probs))] + + if probs.size == 0: + return [] + + def threshold_for(class_name): + if thresholds: + return thresholds.get( + class_name, thresholds.get("default", default_threshold) + ) + return default_threshold + + order = np.argsort(probs)[::-1] + top1 = int(order[0]) + top1_conf = float(probs[top1]) + top2_conf = float(probs[order[1]]) if order.size > 1 else 0.0 + + if top1_conf < threshold_for(names[top1]): + return [] + if (top1_conf - top2_conf) < margin: + return [] + + try: + limit = max(1, int(top_k)) + except (TypeError, ValueError): + limit = 1 + + accepted = [(names[top1], top1_conf)] + for idx in order[1:limit]: + idx = int(idx) + conf = float(probs[idx]) + if conf >= threshold_for(names[idx]): + accepted.append((names[idx], conf)) + return accepted + + +# ------------------------------- aggregation ------------------------------ # + +def wd_aggregate(frame_results): + """Per-label max probability across frames (thresholds already applied).""" + best = {} + for res in frame_results: + for name, score in res: + if name not in best or score > best[name]: + best[name] = float(score) + return sorted(best.items(), key=lambda kv: kv[1], reverse=True) + + +def yolo_aggregate(frame_results, frame_count=None, multi=True, + min_votes=YOLO_MIN_VOTES_DEFAULT, + share=YOLO_SHARE_DEFAULT): + """Aggregate single-label frame results into scene-level tag list. + + The classifier accepts at most one class per frame, but a video can contain + several acts over time, so the accepted frames are pooled: + + * ``multi=True`` (default): every class accepted on at least + ``max(min_votes, ceil(share * accepted_frames))`` frames is returned, so a + video with several distinct acts yields several tags. + * ``multi=False``: the historical single majority-vote winner. + + Ties are broken by higher mean confidence. + """ + counts = {} + sums = {} + accepted_frames = 0 + for res in frame_results: + if not res: + continue + accepted_frames += 1 + for name, score in res: + counts[name] = counts.get(name, 0) + 1 + sums[name] = sums.get(name, 0.0) + float(score) + + if not counts: + return [] + + def _winner(): + best = max(counts, key=lambda n: (counts[n], sums[n] / counts[n])) + return [(best, sums[best] / counts[best])] + + if not multi: + return _winner() + + threshold = max(1, int(min_votes)) + if share and accepted_frames: + threshold = max(threshold, int(math.ceil(share * accepted_frames))) + + kept = [ + (name, sums[name] / counts[name]) + for name, count in counts.items() + if count >= threshold + ] + if not kept: + # Never drop everything: keep the strongest candidate. + return _winner() + + kept.sort(key=lambda item: (counts[item[0]], item[1]), reverse=True) + return kept + + +# ========================================================================== # +# 4. IMAGE PIPELINE +# ========================================================================== # diff --git a/plugins/tagAI/tagai_core/preprocess.py b/plugins/tagAI/tagai_core/preprocess.py new file mode 100755 index 00000000..6ea1161e --- /dev/null +++ b/plugins/tagAI/tagai_core/preprocess.py @@ -0,0 +1,125 @@ +"""Per-family image preprocessing for the ONNX models.""" + +from tagai_core.config import JOYTAG_SIZE, SEXPOS_RESIZE, SEXPOS_SIZE, YOLO_DET_SIZE +from tagai_core.config import Image +from tagai_core.config import np + + +_IMAGENET_MEAN = None if np is None else np.array( + [0.485, 0.456, 0.406], dtype=np.float32) +_IMAGENET_STD = None if np is None else np.array( + [0.229, 0.224, 0.225], dtype=np.float32) + +def wd_preprocess(img, size): + """WD: aspect-preserving resize of the longest side to ``size``, pasted + centred on a white square, BGR, NHWC ``[1, size, size, 3]``.""" + ratio = size / max(img.size) + new = (int(img.width * ratio), int(img.height * ratio)) + img = img.resize(new, Image.LANCZOS) + + square = Image.new("RGB", (size, size), (255, 255, 255)) + square.paste(img, ((size - new[0]) // 2, (size - new[1]) // 2)) + + arr = np.array(square).astype(np.float32)[:, :, ::-1] # RGB -> BGR + return np.expand_dims(arr, 0) # [1, H, W, 3] + + +def yolo_preprocess(img, size=224): + """YOLO classification: resize shorter side to ``size``, centre-crop to + ``size x size``, RGB/255, CHW ``[1, 3, size, size]``.""" + img = img.convert("RGB") + w, h = img.size + scale = size / min(w, h) + nw = max(1, int(round(w * scale))) + nh = max(1, int(round(h * scale))) + img = img.resize((nw, nh), Image.BILINEAR) + + left = (nw - size) // 2 + top = (nh - size) // 2 + img = img.crop((left, top, left + size, top + size)) + + arr = np.asarray(img, dtype=np.float32) / 255.0 + arr = np.transpose(arr, (2, 0, 1)) # HWC -> CHW + return np.expand_dims(arr, 0) + + +def yolo_preprocess_batch(images, size=224): + """Preprocess a list of PIL images into ``[N, 3, size, size]`` float32.""" + if not images: + return None + return np.concatenate( + [yolo_preprocess(img, size) for img in images], axis=0 + ).astype(np.float32) + + +def yolo_det_preprocess(img, size=YOLO_DET_SIZE): + """YOLO11-NSFW detection: letterbox the image to ``size`` with grey padding. + + Returns ``(arr, meta)`` where ``arr`` is ``[1, 3, size, size]`` in [0, 1] + and ``meta`` records the scale and padding used, so detection boxes can be + mapped back to the original image if ever needed. + """ + img = img.convert("RGB") + w, h = img.size + scale = min(size / w, size / h) + nw = max(1, int(round(w * scale))) + nh = max(1, int(round(h * scale))) + resized = img.resize((nw, nh), Image.BILINEAR) + + canvas = Image.new("RGB", (size, size), (114, 114, 114)) + pad_x = (size - nw) // 2 + pad_y = (size - nh) // 2 + canvas.paste(resized, (pad_x, pad_y)) + + arr = np.asarray(canvas, dtype=np.float32) / 255.0 + arr = np.transpose(arr, (2, 0, 1)) # HWC -> CHW + meta = {"scale": scale, "pad_x": pad_x, "pad_y": pad_y, "size": size} + return np.expand_dims(arr, 0), meta + + +def sexpos_preprocess(img, size=SEXPOS_SIZE, resize=SEXPOS_RESIZE): + """porntech/sex-position: resize the shorter side to ``resize`` (BICUBIC), + centre-crop ``size``, scale to [0, 1], ImageNet-normalise -> CHW.""" + img = img.convert("RGB") + w, h = img.size + scale = resize / min(w, h) + nw = max(1, int(round(w * scale))) + nh = max(1, int(round(h * scale))) + img = img.resize((nw, nh), Image.BICUBIC) + + left = (nw - size) // 2 + top = (nh - size) // 2 + img = img.crop((left, top, left + size, top + size)) + + arr = np.asarray(img, dtype=np.float32) / 255.0 + arr = (arr - _IMAGENET_MEAN) / _IMAGENET_STD + arr = np.transpose(arr, (2, 0, 1)) # HWC -> CHW + return np.expand_dims(arr, 0) + + +# JoyTag normalises with the CLIP statistics on a white-padded square. +_JOYTAG_MEAN = None if np is None else np.array( + [0.48145466, 0.4578275, 0.40821073], dtype=np.float32) +_JOYTAG_STD = None if np is None else np.array( + [0.26862954, 0.26130258, 0.27577711], dtype=np.float32) + + +def joytag_preprocess(img, size=JOYTAG_SIZE): + """JoyTag: white-pad the image to a square, resize to ``size`` (BICUBIC), + scale to [0, 1] and CLIP-normalise -> CHW ``[1, 3, size, size]``.""" + img = img.convert("RGB") + w, h = img.size + max_dim = max(w, h) + square = Image.new("RGB", (max_dim, max_dim), (255, 255, 255)) + square.paste(img, ((max_dim - w) // 2, (max_dim - h) // 2)) + + if max_dim != size: + square = square.resize((size, size), Image.BICUBIC) + + arr = np.asarray(square, dtype=np.float32) / 255.0 + arr = (arr - _JOYTAG_MEAN) / _JOYTAG_STD + arr = np.transpose(arr, (2, 0, 1)) # HWC -> CHW + return np.expand_dims(arr, 0) + + +# ------------------------------ postprocessing ---------------------------- # diff --git a/plugins/tagAI/tagai_core/scenes.py b/plugins/tagAI/tagai_core/scenes.py new file mode 100755 index 00000000..8dfdeb4a --- /dev/null +++ b/plugins/tagAI/tagai_core/scenes.py @@ -0,0 +1,168 @@ +"""Scene GraphQL queries, frame extraction and per-scene processing.""" + +from tagai_core import stash_client, stashlog +from tagai_core.debug import debug_add_item +from tagai_core.frames import compute_frame_times, extract_frames +from tagai_core.inference import infer_frames +from tagai_core.markers import combine_hits, create_scene_markers +from tagai_core.tags import filter_tags +from tagai_core.config import ( + FRAME_COUNT_DEFAULT, FRAME_INTERVAL_DEFAULT, FRAME_WORKERS_DEFAULT, + MAX_FRAMES_DEFAULT, PER_PAGE_VIDEO_DEFAULT, +) +import shutil +import tempfile +import traceback + + +def get_scenes(ai_tagged_id, page=1, per_page=PER_PAGE_VIDEO_DEFAULT): + q = """ +query MyQuery($excludes: [ID!], $page: Int, $per_page: Int) { + findScenes( + scene_filter: { tags: { modifier: INCLUDES_ALL, excludes: $excludes } } + filter: { page: $page, per_page: $per_page } + ) { + scenes { id files { path ... on VideoFile { duration width height } } } + count + } +}""" + res = stash_client.gql(q, {"excludes": [ai_tagged_id], "page": page, "per_page": per_page}) + if "errors" in res: + return [], 0 + data = res.get("data", {}).get("findScenes") or {} + return data.get("scenes") or [], data.get("count") or 0 + + +def get_scenes_force(page=1, per_page=PER_PAGE_VIDEO_DEFAULT): + q = """ +query MyQuery($page: Int, $per_page: Int) { + findScenes(filter: { page: $page, per_page: $per_page }) { + scenes { id files { path ... on VideoFile { duration width height } } } + count + } +}""" + res = stash_client.gql(q, {"page": page, "per_page": per_page}) + if "errors" in res: + return [], 0 + data = res.get("data", {}).get("findScenes") or {} + return data.get("scenes") or [], data.get("count") or 0 + + +def get_scenes_w_tag(flag_tag_id, page=1, per_page=PER_PAGE_VIDEO_DEFAULT): + q = """ +query MyQuery($value: [ID!], $page: Int, $per_page: Int) { + findScenes( + scene_filter: { tags: { modifier: INCLUDES_ALL, value: $value } } + filter: { page: $page, per_page: $per_page } + ) { + scenes { id files { path ... on VideoFile { duration width height } } } + count + } +}""" + res = stash_client.gql(q, {"value": [flag_tag_id], "page": page, "per_page": per_page}) + if "errors" in res: + return [], 0 + data = res.get("data", {}).get("findScenes") or {} + return data.get("scenes") or [], data.get("count") or 0 + + +def get_scene_tags(scene_id): + q = """query($id:ID!){ findScene(id:$id){ tags{id} } }""" + r = stash_client.gql(q, {"id": scene_id}) + return r.get("data", {}).get("findScene", {}).get("tags", []) or [] + + +def apply_scene_tags(scene_id, new_tag_ids, exclude_tag_id=None, ai_tagged_id=None): + """Mirror of :func:`apply_tags` using ``sceneUpdate``.""" + new_tag_ids = [t for t in new_tag_ids if t] + if ai_tagged_id: + new_tag_ids.append(ai_tagged_id) + if not new_tag_ids: + stashlog.log.info("No tags to apply.") + return + + existing = [ + t["id"] + for t in get_scene_tags(scene_id) + if t["id"] != exclude_tag_id + ] + final = list(set(existing + new_tag_ids)) + + q = """mutation($id:ID!, $tags:[ID!]!){ + sceneUpdate(input:{id:$id, tag_ids:$tags}){id} + }""" + stash_client.gql(q, {"id": scene_id, "tags": final}) + + +def process_scene(scene, models, cfg, flag_tag_id, ai_tagged_id, prefix): + """Extract frames, classify, tag and (optionally) mark a single scene.""" + files = scene.get("files") or [] + video = files[0] if files else None + path = video.get("path") if video else None + duration = video.get("duration") if video else None + + if not path or not duration or duration <= 0: + stashlog.log.info(f"Scene {scene.get('id')}: no video file/duration, skipping") + return [] + + tmp = tempfile.mkdtemp(prefix="tagai_") + try: + times = compute_frame_times( + duration, + frame_count=cfg.get("frame_count", FRAME_COUNT_DEFAULT), + frame_interval=cfg.get("frame_interval", FRAME_INTERVAL_DEFAULT), + max_frames=cfg.get("max_frames", MAX_FRAMES_DEFAULT), + ) + extracted = extract_frames( + path, times, tmp, cfg["ffmpeg_path"], + workers=cfg.get("frame_workers", FRAME_WORKERS_DEFAULT), + include_times=True, + ) + if not extracted: + stashlog.log.info(f"Scene {scene.get('id')}: no frames extracted, skipping") + return [] + frames = [frame for frame, _moment in extracted] + frame_times = [moment for _frame, moment in extracted] + stashlog.log.info(f"Scene {scene.get('id')}: classifying {len(frames)} frames") + debug_models = {} if cfg.get("debug") else None + markers_mode = cfg.get("markers") or "off" + hits_out = {} if markers_mode != "off" else None + tags = infer_frames( + frames, models, cfg, debug=debug_models, hits_out=hits_out + ) + finally: + shutil.rmtree(tmp, ignore_errors=True) + + tags = filter_tags(tags, cfg.get("tag_whitelist"), cfg.get("tag_blacklist")) + ids = [tid for tid in (stash_client.ensure_tag(prefix + name) for name, _score in tags) if tid] + apply_scene_tags(scene["id"], ids, flag_tag_id, ai_tagged_id) + + created_markers = [] + if markers_mode != "off" and hits_out is not None: + try: + created_markers = create_scene_markers( + scene, tags, combine_hits(hits_out, len(frames)), + frame_times, duration, cfg, prefix, + ) + except Exception as exc: + stashlog.log.error( + f"Scene {scene.get('id')}: could not create markers: {exc}\n" + f"{traceback.format_exc()}" + ) + + if debug_models is not None: + debug_add_item({ + "id": scene.get("id"), + "path": path, + "duration": round(float(duration), 3), + "frames": len(frames), + "models": list(debug_models.values()), + "applied_tags": [prefix + name for name, _score in tags], + "markers": created_markers, + }) + return tags + + +# ========================================================================== # +# 6. JOB LOOP / PAGINATION +# ========================================================================== # diff --git a/plugins/tagAI/tagai_core/stash_client.py b/plugins/tagAI/tagai_core/stash_client.py new file mode 100755 index 00000000..a8a03dde --- /dev/null +++ b/plugins/tagAI/tagai_core/stash_client.py @@ -0,0 +1,157 @@ +"""Stash GraphQL connection and tag lookup/creation helpers.""" + +from tagai_core import config, stashlog +from tagai_core.config import STASH_URL_DEFAULT +from tagai_core.config import requests + + +STASH_URL = STASH_URL_DEFAULT +API_KEY = None +SESSION_COOKIE = None + + +# ========================================================================== # +# 2. CONNECTION & GRAPHQL +# ========================================================================== # + +def setup_connection(json_input): + """Resolve the GraphQL URL + credentials. + + Prefers ``server_connection`` from stdin (so custom ports/auth work without + editing the code) and falls back to the bundled ``api.key`` file. + """ + global STASH_URL, API_KEY, SESSION_COOKIE + + sc = (json_input or {}).get("server_connection") or {} + scheme = sc.get("Scheme") or "http" + host = sc.get("Host") or "localhost" + port = sc.get("Port") or 9999 + STASH_URL = f"{scheme}://{host}:{port}/graphql" + + SESSION_COOKIE = None + sess = sc.get("SessionCookie") + if isinstance(sess, dict) and sess.get("Value"): + SESSION_COOKIE = sess + elif isinstance(sess, str) and sess: + SESSION_COOKIE = {"Name": "session", "Value": sess} + + # api.key file is the battle-tested path; prefer it when present. + api_key = None + try: + with open(config.API_KEY_PATH, encoding="utf-8") as f: + candidate = f.readline().strip() + if candidate: + api_key = candidate + except OSError: + pass + + if api_key is None and SESSION_COOKIE: + api_key = SESSION_COOKIE.get("Value") + + API_KEY = api_key + return STASH_URL, API_KEY, SESSION_COOKIE + + +def gql(query, variables=None): + headers = {} + if API_KEY: + headers["ApiKey"] = API_KEY + if SESSION_COOKIE and SESSION_COOKIE.get("Value"): + headers["Cookie"] = f"{SESSION_COOKIE['Name']}={SESSION_COOKIE['Value']}" + + res = requests.post( + STASH_URL, + json={"query": query, "variables": variables}, + headers=headers, + ).json() + + if "errors" in res: + stashlog.log.error(f"GraphQL Error:\n{res['errors']}") + return res + + +# --------------------------- Tag cache / ensure --------------------------- # + +# Stash treats tag names as case-insensitively unique, so the cache is keyed +# by the lower-cased name: ``AI: christmas`` and ``AI: Christmas`` are the same +# tag as far as Stash is concerned, and re-creating it fails with +# ``tag with name '...' already exists``. +TAG_CACHE = {} # lower-cased name -> id + + +def _cache_tag(tag_id, name): + if tag_id and name: + TAG_CACHE.setdefault(name.lower(), tag_id) + return tag_id + + +def find_tag_by_name(name): + """Return the id of an existing tag, matching names case-insensitively.""" + q = """query($name:String!){ + findTags(tag_filter:{name:{value:$name, modifier:EQUALS}}){ tags{id name} } + }""" + r = gql(q, {"name": name}) + tags = ((r.get("data") or {}).get("findTags") or {}).get("tags") or [] + if tags: + return _cache_tag(tags[0]["id"], tags[0]["name"]) + return None + + +def load_tag_cache(): + """Load every tag name -> id into ``TAG_CACHE`` (paged defensively).""" + page = 1 + per_page = 1000 + while True: + q = """query($page:Int!,$per:Int!){ + findTags(filter:{page:$page, per_page:$per}){ tags{id name} count } + }""" + r = gql(q, {"page": page, "per": per_page}) + data = (r.get("data") or {}).get("findTags") or {} + tags = data.get("tags") or [] + for t in tags: + _cache_tag(t["id"], t["name"]) + count = data.get("count") or 0 + if not tags or len(tags) < per_page or page * per_page >= count: + break + page += 1 + stashlog.log.info(f"Tag cache loaded: {len(TAG_CACHE)} tags") + + +def ensure_tag(name): + """Return the id of ``name``, creating the tag only if it is missing. + + Case differences do not create a second tag: an existing tag whose name + only differs in capitalisation is reused. Returns ``None`` if the tag + could neither be found nor created. + """ + if not name: + return None + + cached = TAG_CACHE.get(name.lower()) + if cached: + return cached + + existing = find_tag_by_name(name) + if existing: + return existing + + q = """mutation($name:String!){ + tagCreate(input:{name:$name}){id name} + }""" + r = gql(q, {"name": name}) + created = (r.get("data") or {}).get("tagCreate") + if created and created.get("id"): + return _cache_tag(created["id"], created.get("name")) + + # Creation failed (e.g. it already existed): look it up one more time. + existing = find_tag_by_name(name) + if existing: + return existing + + stashlog.log.error(f"Could not create or find tag {name!r}") + return None + + +# ========================================================================== # +# 3. MODEL REGISTRY / STRATEGY LAYER (lazy loading) +# ========================================================================== # diff --git a/plugins/tagAI/tagai_core/stashlog.py b/plugins/tagAI/tagai_core/stashlog.py new file mode 100755 index 00000000..7c349e50 --- /dev/null +++ b/plugins/tagAI/tagai_core/stashlog.py @@ -0,0 +1,42 @@ +"""Stash plugin logging (stashapi.log) with a stderr fallback.""" + +import sys + + +try: + import stashapi.log as log +except ImportError: # pragma: no cover - only hit when deps are missing + class _FallbackLog: + def _emit(self, level, *args): + try: + sys.stderr.write( + f"[tagAI][{level}] " + " ".join(str(a) for a in args) + "\n" + ) + sys.stderr.flush() + except Exception: + pass + + def info(self, *a): + self._emit("info", *a) + + def error(self, *a): + self._emit("error", *a) + + def warning(self, *a): + self._emit("warning", *a) + + def debug(self, *a): + self._emit("debug", *a) + + def progress(self, *a): + self._emit("progress", *a) + + log = _FallbackLog() + + +# --------------------------------------------------------------------------- # +# Defensive heavy-dependency import. numpy / requests / Pillow must be present +# in the interpreter Stash uses to run plugins (which is *not* necessarily the +# one on your shell PATH). Record the failure and report it cleanly from +# ``main`` instead of crashing at import time with a bare traceback. +# --------------------------------------------------------------------------- # diff --git a/plugins/tagAI/tagai_core/tags.py b/plugins/tagAI/tagai_core/tags.py new file mode 100755 index 00000000..8aa378ad --- /dev/null +++ b/plugins/tagAI/tagai_core/tags.py @@ -0,0 +1,61 @@ +"""Tag/marker whitelist-blacklist filtering helpers.""" + + + +def parse_tag_filter(value, prefix="AI: "): + """Normalise a marker whitelist/blacklist -> sorted lower-case tag list. + + Entries are separated by commas or semicolons; whitespace is never a + separator because tag names such as ``cowgirl position`` contain spaces. A + leading ``prefix`` (the plugin's ``tag_prefix``) is stripped, so both + ``AI: cowgirl position`` and ``cowgirl position`` work. The result is a + JSON-serialisable list (it is echoed into the debug report's settings). + """ + if isinstance(value, (list, tuple)): + raw = [str(item) for item in value] + else: + raw = str(value or "").replace(";", ",").split(",") + + names = [] + for item in raw: + name = str(item).strip() + if prefix and name.lower().startswith(prefix.lower()): + name = name[len(prefix):] + name = name.strip().lower() + if name and name not in names: + names.append(name) + return sorted(names) + + +def filter_marker_tags(tags, allow=None, deny=None): + """Apply the marker whitelist/blacklist to a ``[(name, score)]`` list. + + ``allow`` (whitelist) restricts markers to the listed tags when it is not + empty; ``deny`` (blacklist) always wins. Names are matched + case-insensitively. Without any filter the list is returned unchanged. + """ + if not allow and not deny: + return list(tags) + + allowed = {str(name).lower() for name in (allow or [])} + denied = {str(name).lower() for name in (deny or [])} + + kept = [] + for name, score in tags: + lowered = str(name).lower() + if lowered in denied: + continue + if allowed and lowered not in allowed: + continue + kept.append((name, score)) + return kept + + +def filter_tags(tags, allow=None, deny=None): + """Apply the scene/image tag whitelist/blacklist to a ``[(name, score)]``. + + Unlike the marker filters this decides which tags are written to the item + itself; the rules are otherwise identical (blacklist wins, case-insensitive + names, an empty list means no filtering). + """ + return filter_marker_tags(tags, allow, deny) diff --git a/plugins/tagAI/tests/test_tagai.py b/plugins/tagAI/tests/test_tagai.py new file mode 100755 index 00000000..6fd4d09c --- /dev/null +++ b/plugins/tagAI/tests/test_tagai.py @@ -0,0 +1,2205 @@ +"""Unit tests for the TagAI plugin. + +These tests deliberately avoid a real ONNX runtime / ffmpeg / Stash instance: +the model layer is exercised with fake sessions and the network layer with +monkey-patched ``gql``. Run with:: + + python -m unittest discover -s tests -v +""" + +import json +import os +import sys +import tempfile +import unittest +from unittest import mock + +import numpy as np +from PIL import Image + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import subprocess # noqa: E402 +from tagai_core import ( # noqa: E402 + cli as tc_cli, + config as tc_config, + debug as tc_debug, + frames as tc_frames, + images as tc_images, + inference as tc_inference, + jobs as tc_jobs, + markers as tc_markers, + models as tc_models, + postprocess as tc_postprocess, + preprocess as tc_preprocess, + scenes as tc_scenes, + stash_client as tc_stash_client, + stashlog as tc_stashlog, + tags as tc_tags, +) + + +# --------------------------------------------------------------------------- # +# Fakes +# --------------------------------------------------------------------------- # + +class _FakeInput: + def __init__(self, name, shape): + self.name = name + self.shape = shape + + +class _FakeOutput: + def __init__(self, name): + self.name = name + + +class _FakeMeta: + def __init__(self, names=None): + self.custom_metadata_map = {} + if names is not None: + self.custom_metadata_map["names"] = names + + +class _FakeSession: + """Minimal stand-in for onnxruntime.InferenceSession.""" + + def __init__(self, path, providers=None, input_shape=(1, 448, 448, 3), + names=None, output_size=2): + self.path = path + self.providers = providers + self._inputs = [_FakeInput("input", list(input_shape))] + self._outputs = [_FakeOutput("output")] + self._meta = _FakeMeta(names) + self.output_size = output_size + + def get_inputs(self): + return self._inputs + + def get_outputs(self): + return self._outputs + + def get_modelmeta(self): + return self._meta + + def run(self, output_names, feeds): + if not feeds: + return [np.full((1, self.output_size), 1.0 / self.output_size, dtype=np.float32)] + batch = next(iter(feeds.values())) + n = int(np.asarray(batch).shape[0]) if np.ndim(batch) else 1 + return [np.full((n, self.output_size), 1.0 / self.output_size, dtype=np.float32)] + + +def _fake_onnxruntime(session_cls): + module = mock.MagicMock() + module.InferenceSession = session_cls + return module + + +# --------------------------------------------------------------------------- # +# Strategy layer: softmax / probabilities +# --------------------------------------------------------------------------- # + +class TestSoftmax(unittest.TestCase): + def test_sums_to_one(self): + out = tc_postprocess.softmax([2.0, 1.0, 0.1]) + self.assertAlmostEqual(float(out.sum()), 1.0, places=6) + + def test_monotonic(self): + out = tc_postprocess.softmax([3.0, 1.0, 0.0]) + self.assertGreater(out[0], out[1]) + self.assertGreater(out[1], out[2]) + + def test_numerically_stable(self): + out = tc_postprocess.softmax([1000.0, 1000.0]) + self.assertTrue(np.all(np.isfinite(out))) + self.assertAlmostEqual(float(out.sum()), 1.0, places=6) + + +class TestAsProbabilities(unittest.TestCase): + def test_detects_softmaxed_output(self): + probs = np.array([0.5, 0.3, 0.2]) + np.testing.assert_allclose(tc_postprocess._as_probabilities(probs), probs) + + def test_applies_softmax_to_logits(self): + out = tc_postprocess._as_probabilities(np.array([2.0, 1.0, 0.1])) + self.assertAlmostEqual(float(out.sum()), 1.0, places=6) + self.assertGreaterEqual(float(out.min()), 0.0) + + +# --------------------------------------------------------------------------- # +# Preprocessing +# --------------------------------------------------------------------------- # + +class TestPreprocess(unittest.TestCase): + def test_wd_preprocess_shape_and_white_pad(self): + img = Image.new("RGB", (200, 100), (255, 0, 0)) # landscape red + arr = tc_preprocess.wd_preprocess(img, 448) + self.assertEqual(arr.shape, (1, 448, 448, 3)) + self.assertEqual(arr.dtype, np.float32) + + # Corners are the white square padding. + np.testing.assert_allclose(arr[0, 0, 0], [255, 255, 255]) + np.testing.assert_allclose(arr[0, -1, -1], [255, 255, 255]) + + # Centre is BGR of red -> [0, 0, 255]. + centre = arr[0, 224, 224] + self.assertLess(centre[0], 5) + self.assertLess(centre[1], 5) + self.assertGreater(centre[2], 250) + + def test_yolo_preprocess_shape_and_range(self): + img = Image.new("RGB", (400, 200), (10, 20, 30)) + arr = tc_preprocess.yolo_preprocess(img, 224) + self.assertEqual(arr.shape, (1, 3, 224, 224)) + self.assertEqual(arr.dtype, np.float32) + self.assertGreaterEqual(float(arr.min()), 0.0) + self.assertLessEqual(float(arr.max()), 1.0) + + +# --------------------------------------------------------------------------- # +# Postprocessing +# --------------------------------------------------------------------------- # + +class TestWdPostprocess(unittest.TestCase): + def setUp(self): + self.model = {"tags": ["a", "b", "c", "d", "e"], + "general_idx": 0, "char_idx": 3} + + def test_general_slice_excludes_characters(self): + probs = np.array([0.9, 0.4, 0.6, 0.99, 0.99]) + result = tc_postprocess.wd_postprocess(probs, self.model, tag_threshold=0.5, + char_threshold=0.85, tag_characters=False) + names = [n for n, _ in result] + self.assertIn("a", names) + self.assertIn("c", names) + self.assertNotIn("d", names) # character row must not leak into general + self.assertNotIn("b", names) # below threshold + + def test_characters_gated(self): + probs = np.array([0.9, 0.4, 0.6, 0.99, 0.99]) + result = tc_postprocess.wd_postprocess(probs, self.model, tag_threshold=0.5, + char_threshold=0.85, tag_characters=True) + names = [n for n, _ in result] + self.assertIn("d", names) + self.assertIn("e", names) + + def test_char_idx_none_uses_end(self): + model = {"tags": ["a", "b", "c"], "general_idx": 0, "char_idx": None} + result = tc_postprocess.wd_postprocess(np.array([0.9, 0.9, 0.9]), model, + tag_threshold=0.5) + self.assertEqual(len(result), 3) + + +class TestYoloPostprocess(unittest.TestCase): + def setUp(self): + self.model = {"class_names": ["A", "B", "C"]} + + def test_accepts_confident_margin(self): + result = tc_postprocess.yolo_postprocess(np.array([0.9, 0.05, 0.05]), self.model, + margin=0.4, default_threshold=0.5) + self.assertEqual(len(result), 1) + self.assertEqual(result[0][0], "A") + self.assertAlmostEqual(result[0][1], 0.9, places=5) + + def test_rejects_small_margin(self): + result = tc_postprocess.yolo_postprocess(np.array([0.55, 0.4, 0.05]), self.model, + margin=0.4, default_threshold=0.5) + self.assertEqual(result, []) + + def test_rejects_low_confidence(self): + result = tc_postprocess.yolo_postprocess(np.array([0.45, 0.3, 0.25]), self.model, + margin=0.1, default_threshold=0.5) + self.assertEqual(result, []) + + def test_per_class_threshold(self): + thresholds = {"A": 0.95, "default": 0.5} + result = tc_postprocess.yolo_postprocess(np.array([0.9, 0.05, 0.05]), self.model, + thresholds=thresholds, margin=0.4) + self.assertEqual(result, []) + + def test_class_name_fallback_when_mismatch(self): + model = {"class_names": []} + result = tc_postprocess.yolo_postprocess(np.array([0.9, 0.05]), model, + margin=0.4, default_threshold=0.5) + self.assertEqual(result[0][0], "class_0") + + def test_top_k_defaults_to_single_class(self): + result = tc_postprocess.yolo_postprocess(np.array([0.95, 0.55, 0.05]), self.model, + margin=0.2, default_threshold=0.5) + self.assertEqual([name for name, _ in result], ["A"]) + + def test_top_k_two_adds_runner_up(self): + result = tc_postprocess.yolo_postprocess(np.array([0.95, 0.55, 0.05]), self.model, + margin=0.2, default_threshold=0.5, top_k=2) + self.assertEqual([name for name, _ in result], ["A", "B"]) + self.assertAlmostEqual(result[1][1], 0.55, places=5) + + def test_top_k_skips_runner_up_below_threshold(self): + result = tc_postprocess.yolo_postprocess(np.array([0.95, 0.3, 0.05]), self.model, + margin=0.2, default_threshold=0.5, top_k=2) + self.assertEqual([name for name, _ in result], ["A"]) + + def test_top_k_uses_per_class_threshold_for_runner_up(self): + thresholds = {"A": 0.5, "B": 0.9, "default": 0.5} + result = tc_postprocess.yolo_postprocess(np.array([0.95, 0.55, 0.05]), self.model, + thresholds=thresholds, margin=0.2, top_k=2) + self.assertEqual([name for name, _ in result], ["A"]) + + def test_top_k_ignored_when_frame_is_rejected(self): + result = tc_postprocess.yolo_postprocess(np.array([0.6, 0.5, 0.05]), self.model, + margin=0.4, default_threshold=0.5, top_k=3) + self.assertEqual(result, []) + + +# --------------------------------------------------------------------------- # +# Aggregation +# --------------------------------------------------------------------------- # + +class TestAggregation(unittest.TestCase): + def test_wd_aggregate_uses_max(self): + frames = [[("a", 0.6)], [("a", 0.9), ("b", 0.7)]] + result = dict(tc_postprocess.wd_aggregate(frames)) + self.assertAlmostEqual(result["a"], 0.9) + self.assertAlmostEqual(result["b"], 0.7) + + def test_yolo_aggregate_majority(self): + frames = [[("A", 0.9)], [("B", 0.8)], [("A", 0.7)]] + result = tc_postprocess.yolo_aggregate(frames) + self.assertEqual(result[0][0], "A") + + def test_yolo_aggregate_tie_breaks_on_mean_conf(self): + frames = [[("A", 0.6)], [("B", 0.9)]] + result = tc_postprocess.yolo_aggregate(frames) + self.assertEqual(result[0][0], "B") + + def test_yolo_aggregate_empty(self): + self.assertEqual(tc_postprocess.yolo_aggregate([[], [], []]), []) + + def test_yolo_multilabel_keeps_every_repeated_class(self): + frames = [[("A", 0.9)], [("A", 0.8)], [("B", 0.7)], [("B", 0.7)]] + result = tc_postprocess.yolo_aggregate(frames, min_votes=2, share=0.0) + self.assertEqual([n for n, _ in result], ["A", "B"]) + + def test_yolo_multilabel_drops_singletons(self): + frames = [[("A", 0.9)], [("A", 0.8)], [("B", 0.7)]] + result = tc_postprocess.yolo_aggregate(frames, min_votes=2, share=0.0) + self.assertEqual([n for n, _ in result], ["A"]) + + def test_yolo_multilabel_share_threshold(self): + # "A" wins 3 of 5 accepted frames, "B" 2 of 5. + frames = [[("A", 0.9)], [("A", 0.9)], [("A", 0.9)], + [("B", 0.9)], [("B", 0.9)]] + # share 0.5 -> a class must win >= 3 frames, so only "A" survives. + self.assertEqual( + [n for n, _ in tc_postprocess.yolo_aggregate(frames, min_votes=1, share=0.5)], + ["A"]) + # share 0.3 -> 2 frames are enough, so both survive. + self.assertEqual( + [n for n, _ in tc_postprocess.yolo_aggregate(frames, min_votes=1, share=0.3)], + ["A", "B"]) + + def test_yolo_single_winner_mode(self): + frames = [[("A", 0.9)], [("B", 0.9)], [("A", 0.8)]] + result = tc_postprocess.yolo_aggregate(frames, multi=False) + self.assertEqual([n for n, _ in result], ["A"]) + + def test_yolo_falls_back_to_winner_when_nothing_passes(self): + frames = [[("A", 0.6)], [("B", 0.9)]] + result = tc_postprocess.yolo_aggregate(frames, min_votes=5, share=0.0) + self.assertEqual([n for n, _ in result], ["B"]) + + +# --------------------------------------------------------------------------- # +# WD label CSV parsing +# --------------------------------------------------------------------------- # + +class TestLoadTags(unittest.TestCase): + def _write_csv(self, rows): + fd, path = tempfile.mkstemp(suffix=".csv") + os.close(fd) + with open(path, "w", encoding="utf-8") as f: + f.write("tag_id,name,category,count\n") + for r in rows: + f.write(",".join(str(x) for x in r) + "\n") + self.addCleanup(os.remove, path) + return path + + def test_boundary_detection(self): + path = self._write_csv([ + (1, "rating_x", 9, 1), + (2, "general_a", 0, 1), + (3, "general_b", 0, 1), + (4, "char_c", 4, 1), + ]) + tags, general_idx, char_idx = tc_models.load_tags(path) + self.assertEqual(tags, ["rating x", "general a", "general b", "char c"]) + self.assertEqual(general_idx, 1) + self.assertEqual(char_idx, 3) + + def test_missing_general_raises(self): + path = self._write_csv([(1, "char_c", 4, 1)]) + with self.assertRaises(RuntimeError): + tc_models.load_tags(path) + + +# --------------------------------------------------------------------------- # +# YOLO metadata / thresholds +# --------------------------------------------------------------------------- # + +class TestYoloMetadata(unittest.TestCase): + def test_class_names_english_passthrough(self): + session = _FakeSession("x", names="{0: 'Cuckold', 1: 'Anime'}") + self.assertEqual(tc_models.load_yolo_class_names(session), ["Cuckold", "Anime"]) + + def test_class_names_translated_from_chinese(self): + """Metadata names are Chinese; they must be mapped to English.""" + session = _FakeSession("x", names="{0: '接吻', 1: '动漫', 2: '未知'}") + self.assertEqual(tc_models.load_yolo_class_names(session), + ["Kissing", "Anime", "未知"]) + + def test_class_names_fallback(self): + session = _FakeSession("x", names=None) + names = tc_models.load_yolo_class_names(session) + self.assertEqual(len(names), 35) + self.assertEqual(names, tc_config.YOLO_CLASS_NAMES) + + def test_bundled_defaults_cover_every_class(self): + for name in tc_config.YOLO_CLASS_NAMES: + self.assertIn(name, tc_config.YOLO_DEFAULT_THRESHOLDS) + self.assertAlmostEqual(tc_config.YOLO_DEFAULT_THRESHOLDS["Kissing"], + 0.64717263, places=8) + self.assertAlmostEqual(tc_config.YOLO_DEFAULT_THRESHOLDS["Pants"], + 0.96068484, places=8) + + def test_load_thresholds_merges_overrides(self): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(tc_config, "MODEL_DIR", tmp): + with open(os.path.join(tmp, "yolo26n-cls-porn-thresholds.json"), + "w", encoding="utf-8") as f: + # Chinese key from the upstream file + a plain override. + f.write('{"接吻": 0.71, "default": 0.4}') + result = tc_models.load_yolo_thresholds("yolo26n-cls-porn") + # Bundled defaults are still present... + self.assertEqual(result["Pants"], tc_config.YOLO_DEFAULT_THRESHOLDS["Pants"]) + # ...and the file overrides win, with zh keys translated. + self.assertAlmostEqual(result["Kissing"], 0.71) + self.assertAlmostEqual(result["default"], 0.4) + + def test_load_thresholds_defaults_without_file(self): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(tc_config, "MODEL_DIR", tmp): + result = tc_models.load_yolo_thresholds("yolo26n-cls-porn") + self.assertEqual(result, tc_config.YOLO_DEFAULT_THRESHOLDS) + + def test_load_thresholds_none_for_wd(self): + self.assertIsNone(tc_models.load_yolo_thresholds("wd-v1-4-moat-tagger-v2")) + + +# --------------------------------------------------------------------------- # +# Model registry +# --------------------------------------------------------------------------- # + +class TestModelRegistry(unittest.TestCase): + def test_unknown_model_raises(self): + with self.assertRaises(ValueError): + tc_models.get_model("does-not-exist") + + def test_missing_file_raises_clear_error(self): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(tc_config, "MODEL_DIR", tmp): + with self.assertRaises(FileNotFoundError) as ctx: + tc_models.get_model("wd-v1-4-moat-tagger-v2") + self.assertIn("README", str(ctx.exception)) + + def test_loads_wd_model_and_caches(self): + with tempfile.TemporaryDirectory() as tmp: + with open(os.path.join(tmp, "wd-v1-4-moat-tagger-v2.onnx"), "w") as f: + f.write("x") + with open(os.path.join(tmp, "wd-v1-4-moat-tagger-v2.csv"), "w", + encoding="utf-8") as f: + f.write("tag_id,name,category,count\n") + f.write("1,general_a,0,1\n") + f.write("2,char_c,4,1\n") + + fake = _fake_onnxruntime(_FakeSession) + tc_models._SESSION_CACHE.clear() + self.addCleanup(tc_models._SESSION_CACHE.clear) + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.dict(sys.modules, {"onnxruntime": fake}): + model = tc_models.get_model("wd-v1-4-moat-tagger-v2") + self.assertEqual(model["family"], "wd") + self.assertEqual(model["size"], 448) + self.assertEqual(model["tags"], ["general a", "char c"]) + # cached + self.assertIs(tc_models.get_model("wd-v1-4-moat-tagger-v2"), model) + + def test_softmax_applied_for_yolo(self): + with tempfile.TemporaryDirectory() as tmp: + with open(os.path.join(tmp, "yolo26n-cls-porn.onnx"), "w") as f: + f.write("x") + + fake = _fake_onnxruntime( + lambda path, providers=None: _FakeSession( + path, providers, input_shape=(1, 3, 224, 224), + names="{0: 'Cuckold', 1: 'Anime'},", output_size=2, + ) + ) + tc_models._SESSION_CACHE.clear() + self.addCleanup(tc_models._SESSION_CACHE.clear) + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.dict(sys.modules, {"onnxruntime": fake}): + model = tc_models.get_model("yolo26n-cls-porn") + self.assertEqual(model["family"], "yolo") + self.assertEqual(model["size"], 224) + + +class _FakeResponse: + """Minimal stand-in for a requests streaming response.""" + + def __init__(self, status=200, body=b"", headers=None): + self.status_code = status + self._body = body + self.headers = headers if headers is not None else { + "Content-Length": str(len(body)) + } + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def raise_for_status(self): + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") + + def iter_content(self, chunk_size=1): + yield self._body + + +class TestModelDownload(unittest.TestCase): + def setUp(self): + # Keep the environment token-free by default. + self._env = mock.patch.dict( + os.environ, + {"HF_TOKEN": "", "HUGGING_FACE_HUB_TOKEN": "", "HUGGINGFACE_TOKEN": ""}, + clear=False, + ) + self._env.start() + self.addCleanup(self._env.stop) + + self._log = mock.patch.object(tc_stashlog, "log", mock.MagicMock()) + self._log.start() + self.addCleanup(self._log.stop) + + def test_token_precedence_explicit(self): + self.assertEqual(tc_models.resolve_hf_token("explicit"), "explicit") + + def test_token_from_env(self): + with mock.patch.dict(os.environ, {"HF_TOKEN": "abc123"}): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(tc_config, "PLUGIN_DIR", tmp): + self.assertEqual(tc_models.resolve_hf_token(), "abc123") + + def test_token_from_file(self): + with tempfile.TemporaryDirectory() as tmp: + with open(os.path.join(tmp, "hf_token"), "w") as f: + f.write("filetoken\n") + with mock.patch.object(tc_config, "PLUGIN_DIR", tmp): + self.assertEqual(tc_models.resolve_hf_token(), "filetoken") + + def test_token_absent_returns_none(self): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(tc_config, "PLUGIN_DIR", tmp): + self.assertIsNone(tc_models.resolve_hf_token()) + + def test_hf_url(self): + self.assertEqual( + tc_models._hf_url("owner/repo", "weights/best.onnx"), + f"{tc_config.HF_ENDPOINT}/owner/repo/resolve/main/weights/best.onnx", + ) + + def test_ensure_skips_when_present(self): + with tempfile.TemporaryDirectory() as tmp: + for local, _ in tc_config.MODELS["wd-v1-4-moat-tagger-v2"]["files"]: + open(os.path.join(tmp, local), "w").close() + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.object(tc_config.requests, "get") as get: + self.assertEqual( + tc_models.ensure_model_files("wd-v1-4-moat-tagger-v2"), [] + ) + get.assert_not_called() + + def test_ensure_downloads_missing_files(self): + with tempfile.TemporaryDirectory() as tmp: + resp = _FakeResponse(body=b"HELLO") + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.object(tc_config.requests, "get", + return_value=resp) as get: + out = tc_models.ensure_model_files("wd-v1-4-moat-tagger-v2") + self.assertEqual(len(out), 2) + self.assertEqual(get.call_count, 2) + for local, _ in tc_config.MODELS["wd-v1-4-moat-tagger-v2"]["files"]: + with open(os.path.join(tmp, local), "rb") as f: + self.assertEqual(f.read(), b"HELLO") + + def test_ensure_download_disabled_raises(self): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(tc_config, "MODEL_DIR", tmp): + with self.assertRaises(FileNotFoundError) as ctx: + tc_models.ensure_model_files( + "wd-v1-4-moat-tagger-v2", allow_download=False + ) + self.assertIn("README", str(ctx.exception)) + + def test_gated_permission_error_is_clear(self): + with tempfile.TemporaryDirectory() as tmp: + resp = _FakeResponse(status=403) + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.object(tc_config.requests, "get", return_value=resp): + with self.assertRaises(FileNotFoundError) as ctx: + tc_models.ensure_model_files("yolo26n-cls-porn", + token="tok") + message = str(ctx.exception) + self.assertIn("token", message.lower()) + self.assertIn("README", message) + + def test_unknown_model_raises(self): + with self.assertRaises(ValueError): + tc_models.ensure_model_files("nope") + + def test_optional_file_downloaded_best_effort(self): + with tempfile.TemporaryDirectory() as tmp: + # Required ONNX present; only the optional thresholds file missing. + open(os.path.join(tmp, "yolo26n-cls-porn.onnx"), "w").close() + resp = _FakeResponse(body=b'{"a": 0.5}') + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.object(tc_config.requests, "get", + return_value=resp) as get: + out = tc_models.ensure_model_files("yolo26n-cls-porn", token="tok") + self.assertEqual(len(out), 1) + self.assertTrue(out[0].endswith("yolo26n-cls-porn-thresholds.json")) + self.assertEqual(get.call_count, 1) + + def test_optional_file_failure_is_not_fatal(self): + with tempfile.TemporaryDirectory() as tmp: + open(os.path.join(tmp, "yolo26n-cls-porn.onnx"), "w").close() + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.object(tc_config.requests, "get", + side_effect=RuntimeError("boom")): + self.assertEqual( + tc_models.ensure_model_files("yolo26n-cls-porn", token="tok"), + []) + + +# --------------------------------------------------------------------------- # +# Connection resolution +# --------------------------------------------------------------------------- # + +class TestConnection(unittest.TestCase): + def setUp(self): + self._saved = (tc_stash_client.STASH_URL, tc_stash_client.API_KEY, tc_stash_client.SESSION_COOKIE, + tc_config.API_KEY_PATH) + + def tearDown(self): + (tc_stash_client.STASH_URL, tc_stash_client.API_KEY, tc_stash_client.SESSION_COOKIE, + tc_config.API_KEY_PATH) = self._saved + + def test_from_server_connection(self): + tc_config.API_KEY_PATH = "/nonexistent/api.key" + url, key, cookie = tc_stash_client.setup_connection({ + "server_connection": { + "Scheme": "https", "Host": "stash.local", "Port": 8443, + "SessionCookie": "abc123", + } + }) + self.assertEqual(url, "https://stash.local:8443/graphql") + self.assertEqual(key, "abc123") + self.assertEqual(cookie, {"Name": "session", "Value": "abc123"}) + + def test_cookie_dict(self): + tc_config.API_KEY_PATH = "/nonexistent/api.key" + _url, key, cookie = tc_stash_client.setup_connection({ + "server_connection": { + "SessionCookie": {"Name": "session", "Value": "xyz"}, + } + }) + self.assertEqual(key, "xyz") + self.assertEqual(cookie["Value"], "xyz") + + def test_falls_back_to_api_key_file(self): + with tempfile.NamedTemporaryFile("w", delete=False) as f: + f.write("file-key\n") + path = f.name + self.addCleanup(os.remove, path) + tc_config.API_KEY_PATH = path + _url, key, _cookie = tc_stash_client.setup_connection({}) + self.assertEqual(key, "file-key") + + +# --------------------------------------------------------------------------- # +# GraphQL queries + apply mutations +# --------------------------------------------------------------------------- # + +class TestGraphQLQueries(unittest.TestCase): + def test_get_images_parses_response(self): + captured = {} + + def fake_gql(query, variables=None): + captured.update(variables) + return {"data": {"findImages": { + "images": [{"id": "1"}], "count": 7}}} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + images, count = tc_images.get_images("1357", page=2, per_page=10) + self.assertEqual(images, [{"id": "1"}]) + self.assertEqual(count, 7) + self.assertEqual(captured["excludes"], ["1357"]) + self.assertEqual(captured["page"], 2) + self.assertEqual(captured["per_page"], 10) + + def test_get_scenes_video_fields(self): + def fake_gql(query, variables=None): + self.assertIn("... on VideoFile", query) + return {"data": {"findScenes": { + "scenes": [{"id": "3968"}], "count": 1}}} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + scenes, count = tc_scenes.get_scenes("1357") + self.assertEqual(scenes, [{"id": "3968"}]) + self.assertEqual(count, 1) + + def test_apply_tags_merge(self): + calls = [] + + def fake_gql(query, variables=None): + calls.append((query, variables)) + if "findImage" in query: + return {"data": {"findImage": {"tags": [ + {"id": "10"}, {"id": "1355"}, {"id": "20"}]}}} + return {"data": {"imageUpdate": {"id": "1"}}} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + tc_images.apply_tags("1", ["100", "200"], + exclude_tag_id="1355", ai_tagged_id="1357") + + mutation = [c for c in calls if "imageUpdate" in c[0]][0] + self.assertEqual(set(mutation[1]["tags"]), + {"10", "20", "100", "200", "1357"}) + + def test_apply_scene_tags_merge(self): + calls = [] + + def fake_gql(query, variables=None): + calls.append((query, variables)) + if "findScene" in query: + return {"data": {"findScene": {"tags": [{"id": "10"}]}}} + return {"data": {"sceneUpdate": {"id": "1"}}} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + tc_scenes.apply_scene_tags("1", ["100"], + exclude_tag_id="1355", ai_tagged_id="1357") + + mutation = [c for c in calls if "sceneUpdate" in c[0]][0] + self.assertEqual(set(mutation[1]["tags"]), {"10", "100", "1357"}) + + +# --------------------------------------------------------------------------- # +# Tag cache / ensure_tag (case-insensitive tag names) +# --------------------------------------------------------------------------- # + +class TestTagCache(unittest.TestCase): + def setUp(self): + self._saved = dict(tc_stash_client.TAG_CACHE) + tc_stash_client.TAG_CACHE = {} + self._log = mock.patch.object(tc_stashlog, "log", mock.MagicMock()) + self._log.start() + + def tearDown(self): + self._log.stop() + tc_stash_client.TAG_CACHE = self._saved + + def test_load_tag_cache_is_case_insensitive(self): + def fake_gql(query, variables=None): + return {"data": {"findTags": { + "tags": [{"id": "3290", "name": "AI: christmas"}], + "count": 1}}} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + tc_stash_client.load_tag_cache() + self.assertEqual(tc_stash_client.TAG_CACHE["ai: christmas"], "3290") + + def test_ensure_tag_uses_cache_without_querying(self): + tc_stash_client.TAG_CACHE["ai: x"] = "5" + + def boom(query, variables=None): + raise AssertionError("gql must not be called on a cache hit") + + with mock.patch.object(tc_stash_client, "gql", boom): + self.assertEqual(tc_stash_client.ensure_tag("AI: X"), "5") + + def test_ensure_tag_reuses_existing_case_variant(self): + def fake_gql(query, variables=None): + if "findTags" in query: + return {"data": {"findTags": { + "tags": [{"id": "3290", "name": "AI: christmas"}]}}} + raise AssertionError("tagCreate must not run for an existing tag") + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + self.assertEqual(tc_stash_client.ensure_tag("AI: Christmas"), "3290") + self.assertEqual(tc_stash_client.TAG_CACHE["ai: christmas"], "3290") + + def test_ensure_tag_creates_when_missing(self): + def fake_gql(query, variables=None): + if "findTags" in query: + return {"data": {"findTags": {"tags": []}}} + return {"data": {"tagCreate": {"id": "9", "name": "AI: New"}}} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + self.assertEqual(tc_stash_client.ensure_tag("AI: New"), "9") + self.assertEqual(tc_stash_client.TAG_CACHE["ai: new"], "9") + + def test_ensure_tag_recovers_when_create_fails(self): + state = {"lookups": 0} + + def fake_gql(query, variables=None): + if "findTags" in query: + state["lookups"] += 1 + # Not found before the create, found afterwards. + tags = [] if state["lookups"] == 1 else [ + {"id": "7", "name": "AI: Kissing"}] + return {"data": {"findTags": {"tags": tags}}} + return {"errors": [{"message": "already exists"}], "data": None} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + self.assertEqual(tc_stash_client.ensure_tag("AI: Kissing"), "7") + self.assertEqual(state["lookups"], 2) + + def test_ensure_tag_returns_none_when_unresolvable(self): + def fake_gql(query, variables=None): + if "findTags" in query: + return {"data": {"findTags": {"tags": []}}} + return {"errors": [{"message": "boom"}], "data": None} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + self.assertIsNone(tc_stash_client.ensure_tag("AI: Nope")) + tc_stashlog.log.error.assert_called() + + def test_ensure_tag_blank_name(self): + self.assertIsNone(tc_stash_client.ensure_tag("")) + + +# --------------------------------------------------------------------------- # +# Pagination loop +# --------------------------------------------------------------------------- # + +class TestIterateJob(unittest.TestCase): + def test_tagall_shrinking_set_terminates(self): + state = {"ids": [1, 2, 3, 4, 5], "fetches": 0} + + def fetch(page): + state["fetches"] += 1 + items = [{"id": i} for i in state["ids"]] + return items, len(items) + + def process(batch): + for it in batch: + state["ids"].remove(it["id"]) + + processed = tc_jobs.iterate_job(fetch, process, "TagAll", 2) + self.assertEqual(processed, {1, 2, 3, 4, 5}) + self.assertEqual(state["ids"], []) + self.assertLessEqual(state["fetches"], 10) # no infinite loop + + def test_tagall_stops_when_nothing_new(self): + fetches = {"n": 0} + + def fetch(page): + fetches["n"] += 1 + return [{"id": 1}, {"id": 2}], 2 + + processed = tc_jobs.iterate_job(fetch, lambda b: None, "TagAll", 2) + self.assertEqual(processed, {1, 2}) + self.assertEqual(fetches["n"], 2) # second fetch finds no new ids + + def test_forcetag_advances_pages(self): + pages = {1: [{"id": 1}, {"id": 2}], 2: [{"id": 3}], 3: []} + seen_pages = [] + + def fetch(page): + seen_pages.append(page) + return pages.get(page, []), 3 + + processed = tc_jobs.iterate_job(fetch, lambda b: None, "ForceTag", 2) + self.assertEqual(processed, {1, 2, 3}) + self.assertIn(2, seen_pages) + self.assertNotIn(1, [seen_pages[i] for i in range(len(seen_pages)) + if i > 0]) # page advances, never resets + + def test_progress_callback(self): + progress = [] + pages = {1: [{"id": 1}], 2: []} + + def fetch(page): + return pages.get(page, []), 1 + + tc_jobs.iterate_job(fetch, lambda b: None, "ForceTag", 1, + on_progress=progress.append) + self.assertTrue(progress) + self.assertLessEqual(max(progress), 1.0) + + +# --------------------------------------------------------------------------- # +# Frame extraction +# --------------------------------------------------------------------------- # + +class TestComputeFrameTimes(unittest.TestCase): + def test_zero_or_missing_duration(self): + self.assertEqual(tc_frames.compute_frame_times(0), []) + self.assertEqual(tc_frames.compute_frame_times(None), []) + + def test_frame_count_mode(self): + """frame_interval <= 0 falls back to an evenly spread frame_count.""" + times = tc_frames.compute_frame_times(100.0, frame_count=4, frame_interval=0) + self.assertEqual(len(times), 4) + self.assertGreaterEqual(min(times), 10.0) + self.assertLessEqual(max(times), 90.0) + # Evenly spread across the band. + self.assertAlmostEqual(times[1] - times[0], times[2] - times[1], places=5) + + def test_interval_mode_counts_frames(self): + """~one frame per interval, sampled across the middle 80% band.""" + times = tc_frames.compute_frame_times(100.0, frame_interval=10.0, + max_frames=1000) + # span = 80s, 80 // 10 + 1 == 9 samples. + self.assertEqual(len(times), 9) + self.assertGreaterEqual(min(times), 10.0) + self.assertLessEqual(max(times), 90.0) + + def test_interval_mode_honours_max_frames(self): + times = tc_frames.compute_frame_times(3600.0, frame_interval=1.0, + max_frames=64) + self.assertEqual(len(times), 64) + + def test_short_video_yields_single_frame(self): + times = tc_frames.compute_frame_times(5.0, frame_interval=10.0) + self.assertEqual(len(times), 1) + self.assertGreater(times[0], 0.0) + self.assertLess(times[0], 5.0) + + +class TestExtractFrames(unittest.TestCase): + def test_no_video_or_times_returns_empty(self): + with tempfile.TemporaryDirectory() as tmp: + self.assertEqual(tc_frames.extract_frames("", [1.0], tmp), []) + self.assertEqual(tc_frames.extract_frames("x.mp4", None, tmp), []) + self.assertEqual(tc_frames.extract_frames("x.mp4", [], tmp), []) + + def test_missing_video_returns_empty(self): + with tempfile.TemporaryDirectory() as tmp: + self.assertEqual( + tc_frames.extract_frames("/no/such/video.mp4", [50.0], tmp), []) + + def test_missing_ffmpeg_returns_empty(self): + with tempfile.TemporaryDirectory() as tmp: + video = os.path.join(tmp, "v.mp4") + open(video, "w").close() + with mock.patch.object(subprocess, "run", + side_effect=FileNotFoundError): + frames = tc_frames.extract_frames(video, [30.0, 60.0], tmp, + ffmpeg_path="nope") + self.assertEqual(frames, []) + + def test_extracts_one_frame_per_time(self): + captured = [] + with tempfile.TemporaryDirectory() as tmp: + video = os.path.join(tmp, "v.mp4") + open(video, "w").close() + times = tc_frames.compute_frame_times(100.0, frame_count=4, + frame_interval=0) + + class _Proc: + returncode = 0 + stderr = "" + + def fake_run(cmd, capture_output=True, text=True): + captured.append(cmd) + # Emulate ffmpeg writing the output frame. + out = cmd[-1] + open(out, "w").close() + return _Proc() + + with mock.patch.object(subprocess, "run", fake_run): + frames = tc_frames.extract_frames(video, times, tmp) + + self.assertEqual(len(frames), 4) + ss_times = [float(cmd[cmd.index("-ss") + 1]) for cmd in captured] + self.assertGreaterEqual(min(ss_times), 10.0) + self.assertLessEqual(max(ss_times), 90.0) + + def test_include_times_returns_timestamps(self): + with tempfile.TemporaryDirectory() as tmp: + video = os.path.join(tmp, "v.mp4") + open(video, "w").close() + times = tc_frames.compute_frame_times(100.0, frame_count=4, + frame_interval=0) + + class _Proc: + returncode = 0 + stderr = "" + + def fake_run(cmd, capture_output=True, text=True): + open(cmd[-1], "w").close() + return _Proc() + + with mock.patch.object(subprocess, "run", fake_run): + frames = tc_frames.extract_frames(video, times, tmp, + include_times=True) + + self.assertEqual(len(frames), 4) + self.assertEqual([moment for _path, moment in frames], list(times)) + self.assertTrue(all(isinstance(path, str) + for path, _moment in frames)) + + def test_parallel_extraction_returns_all_frames(self): + captured = [] + with tempfile.TemporaryDirectory() as tmp: + video = os.path.join(tmp, "v.mp4") + open(video, "w").close() + times = tc_frames.compute_frame_times(300.0, frame_count=12, + frame_interval=0) + + class _Proc: + returncode = 0 + stderr = "" + + def fake_run(cmd, capture_output=True, text=True): + captured.append(cmd) + open(cmd[-1], "w").close() + return _Proc() + + with mock.patch.object(subprocess, "run", fake_run): + frames = tc_frames.extract_frames(video, times, tmp, workers=4) + + self.assertEqual(len(frames), 12) + self.assertEqual(len(captured), 12) + # Frames must come back ordered by index regardless of completion + # order. + self.assertEqual( + frames, + [os.path.join(tmp, f"frame_{i:03d}.jpg") for i in range(12)], + ) + + +# --------------------------------------------------------------------------- # +# Video inference batching +# --------------------------------------------------------------------------- # + +class TestInferFrames(unittest.TestCase): + def _write_frames(self, tmp, count): + paths = [] + for i in range(count): + path = os.path.join(tmp, f"frame_{i:03d}.jpg") + Image.new("RGB", (64, 64), (i * 10 % 255, 0, 0)).save(path) + paths.append(path) + return paths + + def _yolo_model(self, session): + return { + "family": "yolo", + "size": 224, + "session": session, + "input_name": "input", + "output_name": "output", + "class_names": ["AI Generated", "Anime"], + "thresholds": None, + } + + def test_yolo_runs_batched(self): + with tempfile.TemporaryDirectory() as tmp: + paths = self._write_frames(tmp, 5) + session = mock.MagicMock() + session.run.return_value = [ + np.array([[1.0, 0.0]], dtype=np.float32) for _ in range(4) + ] + cfg = {"yolo_margin": 0.0, "yolo_threshold": 0.5, + "yolo_multi": True, "yolo_min_votes": 1, "yolo_share": 0.0, + "infer_batch_size": 4} + + results = tc_inference.infer_frames(paths, self._yolo_model(session), cfg) + + # 5 frames in batches of 4 -> two session calls (4 + 1). + self.assertEqual(session.run.call_count, 2) + batch_sizes = [ + call.args[1]["input"].shape[0] for call in session.run.call_args_list + ] + self.assertEqual(batch_sizes, [4, 1]) + self.assertEqual(results[0][0], "AI Generated") + + def test_yolo_batch_size_one_still_works(self): + with tempfile.TemporaryDirectory() as tmp: + paths = self._write_frames(tmp, 3) + session = mock.MagicMock() + session.run.return_value = [ + np.array([[0.7, 0.2]], dtype=np.float32) for _ in range(3) + ] + cfg = {"yolo_margin": 0.0, "yolo_threshold": 0.5, + "yolo_multi": True, "yolo_min_votes": 1, "yolo_share": 0.0, + "infer_batch_size": 1} + + results = tc_inference.infer_frames(paths, self._yolo_model(session), cfg) + + self.assertEqual(session.run.call_count, 3) + self.assertEqual(results, [("AI Generated", results[0][1])]) + + def test_unreadable_frames_are_skipped(self): + with tempfile.TemporaryDirectory() as tmp: + paths = self._write_frames(tmp, 2) + paths.insert(1, os.path.join(tmp, "missing.jpg")) + session = mock.MagicMock() + session.run.return_value = [ + np.array([[1.0, 0.0]], dtype=np.float32) for _ in range(2) + ] + cfg = {"yolo_margin": 0.0, "yolo_threshold": 0.5, + "yolo_multi": True, "yolo_min_votes": 1, "yolo_share": 0.0, + "infer_batch_size": 8} + + with mock.patch.object(tc_stashlog, "log", mock.MagicMock()): + results = tc_inference.infer_frames(paths, self._yolo_model(session), cfg) + + self.assertEqual(session.run.call_count, 1) + self.assertEqual(session.run.call_args.args[1]["input"].shape[0], 2) + self.assertEqual(results[0][0], "AI Generated") + + +# --------------------------------------------------------------------------- # +# Config parsing +# --------------------------------------------------------------------------- # + +class TestJoyTag(unittest.TestCase): + def test_load_joytag_tags(self): + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "top_tags.txt") + with open(path, "w", encoding="utf-8") as f: + f.write("alpha\nbeta\n\n gamma \n") + self.assertEqual( + tc_models.load_joytag_tags(path), ["alpha", "beta", "gamma"] + ) + + def test_sigmoid_passthrough_when_already_probabilities(self): + arr = np.array([0.1, 0.9, 0.5], dtype=np.float32) + out = tc_postprocess._as_sigmoid_probabilities(arr) + np.testing.assert_allclose(out, arr) + + def test_sigmoid_applied_to_logits(self): + arr = np.array([-1.0, 0.0, 1.0], dtype=np.float32) + out = tc_postprocess._as_sigmoid_probabilities(arr) + expected = 1.0 / (1.0 + np.exp(-arr)) + np.testing.assert_allclose(out, expected, rtol=1e-5) + self.assertGreater(float(out.min()), 0.0) + self.assertLess(float(out.max()), 1.0) + + def test_preprocess_shape_and_normalisation(self): + img = Image.new("RGB", (100, 50), (255, 255, 255)) + arr = tc_preprocess.joytag_preprocess(img) + self.assertEqual(arr.shape, (1, 3, 448, 448)) + # A white pixel normalises to (1 - mean) / std for every channel. + self.assertAlmostEqual( + float(arr[0, 0, 0, 0]), (1.0 - 0.48145466) / 0.26862954, places=4 + ) + + def test_postprocess_keeps_tags_above_threshold(self): + model = {"tags": ["a", "b", "c"]} + probs = np.array([0.9, 0.1, 0.5], dtype=np.float32) + hits = tc_postprocess.joytag_postprocess(probs, model) + self.assertEqual([name for name, _ in hits], ["a", "c"]) + self.assertAlmostEqual(hits[0][1], 0.9, places=5) + self.assertEqual(tc_postprocess.joytag_postprocess(probs, model, threshold=0.95), []) + + def test_postprocess_name_fallback(self): + model = {"tags": ["a"]} + probs = np.array([0.9, 0.8], dtype=np.float32) + hits = tc_postprocess.joytag_postprocess(probs, model) + self.assertEqual([name for name, _ in hits], ["tag_0", "tag_1"]) + + def test_registry_loads_joytag(self): + with tempfile.TemporaryDirectory() as tmp: + with open(os.path.join(tmp, "joytag.onnx"), "w") as f: + f.write("x") + with open(os.path.join(tmp, "joytag-top_tags.txt"), "w", + encoding="utf-8") as f: + f.write("alpha\nbeta\n") + + fake = _fake_onnxruntime(_FakeSession) + tc_models._SESSION_CACHE.clear() + self.addCleanup(tc_models._SESSION_CACHE.clear) + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.dict(sys.modules, {"onnxruntime": fake}): + model = tc_models.get_model("joytag") + self.assertEqual(model["family"], "joytag") + self.assertEqual(model["size"], tc_config.JOYTAG_SIZE) + self.assertEqual(model["tags"], ["alpha", "beta"]) + + def test_registry_loads_wd_v3(self): + with tempfile.TemporaryDirectory() as tmp: + with open(os.path.join(tmp, "wd-swinv2-tagger-v3.onnx"), "w") as f: + f.write("x") + with open(os.path.join(tmp, "wd-swinv2-tagger-v3.csv"), "w", + encoding="utf-8") as f: + f.write("tag_id,name,category,count\n") + f.write("1,v3_general,0,1\n") + + fake = _fake_onnxruntime(_FakeSession) + tc_models._SESSION_CACHE.clear() + self.addCleanup(tc_models._SESSION_CACHE.clear) + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.dict(sys.modules, {"onnxruntime": fake}): + model = tc_models.get_model("wd-swinv2-tagger-v3") + self.assertEqual(model["family"], "wd") + self.assertEqual(model["tags"], ["v3 general"]) + + def test_joytag_threshold_default_and_override(self): + self.assertEqual( + tc_cli.build_cfg({})["joytag_threshold"], + tc_config.JOYTAG_THRESHOLD_DEFAULT, + ) + self.assertAlmostEqual( + tc_cli.build_cfg({"joytag_threshold": "0.6"})["joytag_threshold"], + 0.6, + ) + + +class TestYoloDet(unittest.TestCase): + def test_class_names_from_metadata(self): + session = _FakeSession("x", names="{0: 'person', 1: 'blowjob'}") + self.assertEqual(tc_models.load_det_class_names(session), ["person", "blowjob"]) + + def test_class_names_fallback(self): + self.assertEqual( + tc_models.load_det_class_names(_FakeSession("x")), + tc_config.YOLO_DET_CLASS_NAMES, + ) + + def test_preprocess_letterboxes_to_square(self): + img = Image.new("RGB", (1280, 720), (255, 255, 255)) + arr, meta = tc_preprocess.yolo_det_preprocess(img) + self.assertEqual(arr.shape, (1, 3, 640, 640)) + self.assertEqual(meta["size"], 640) + self.assertAlmostEqual(meta["scale"], 0.5, places=6) + self.assertEqual(meta["pad_x"], 0) + self.assertEqual(meta["pad_y"], 140) + # Padding is grey (114) and the image content is untouched white. + self.assertAlmostEqual(float(arr[0, 0, 0, 0]), 114 / 255.0, places=5) + self.assertAlmostEqual(float(arr[0, 0, 320, 320]), 1.0, places=5) + + def test_nms_suppresses_overlapping_boxes(self): + boxes = [[0, 0, 10, 10], [1, 1, 11, 11], [100, 100, 110, 110]] + scores = [0.9, 0.8, 0.7] + self.assertEqual(tc_postprocess._nms(boxes, scores, 0.45), [0, 2]) + + @staticmethod + def _det_output(entries, num_classes=10, anchors=8400): + """Build a synthetic ``[1, 4 + num_classes, anchors]`` head output.""" + out = np.zeros((1, 4 + num_classes, anchors), dtype=np.float32) + for anchor, (box, class_id, conf) in enumerate(entries): + out[0, :4, anchor] = box + out[0, 4 + class_id, anchor] = conf + return out + + def test_postprocess_decodes_a_detection(self): + model = {"class_names": tc_config.YOLO_DET_CLASS_NAMES} + out = self._det_output([([320.0, 320.0, 100.0, 100.0], 8, 0.9)]) + hits = tc_postprocess.yolo_det_postprocess(out, model) + self.assertEqual(len(hits), 1) + self.assertEqual(hits[0][0], "blowjob") + self.assertAlmostEqual(hits[0][1], 0.9, places=5) + + def test_postprocess_respects_conf_threshold(self): + model = {"class_names": tc_config.YOLO_DET_CLASS_NAMES} + out = self._det_output([([320.0, 320.0, 100.0, 100.0], 8, 0.1)]) + self.assertEqual(tc_postprocess.yolo_det_postprocess(out, model), []) + self.assertEqual( + len(tc_postprocess.yolo_det_postprocess(out, model, conf_threshold=0.05)), 1 + ) + + def test_postprocess_suppresses_duplicate_boxes(self): + model = {"class_names": tc_config.YOLO_DET_CLASS_NAMES} + out = self._det_output([ + ([320.0, 320.0, 100.0, 100.0], 8, 0.9), + ([322.0, 320.0, 100.0, 100.0], 7, 0.5), + ]) + self.assertEqual([name for name, _ in tc_postprocess.yolo_det_postprocess(out, model)], + ["blowjob"]) + + def test_postprocess_accepts_transposed_layout(self): + model = {"class_names": ["a", "b"]} + out = np.zeros((8400, 6), dtype=np.float32) + out[0, :4] = [10.0, 10.0, 20.0, 20.0] + out[0, 5] = 0.8 + hits = tc_postprocess.yolo_det_postprocess(out, model) + self.assertEqual(len(hits), 1) + self.assertEqual(hits[0][0], "b") + + def test_registry_loads_detector(self): + with tempfile.TemporaryDirectory() as tmp: + with open(os.path.join(tmp, "yolo11-nsfw.onnx"), "w") as f: + f.write("x") + fake = _fake_onnxruntime( + lambda path, providers=None: _FakeSession( + path, providers, + input_shape=(1, 3, 640, 640), + names="{0: 'person', 9: 'handjob'}", + output_size=14, + ) + ) + tc_models._SESSION_CACHE.clear() + self.addCleanup(tc_models._SESSION_CACHE.clear) + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.dict(sys.modules, {"onnxruntime": fake}): + model = tc_models.get_model("yolo11-nsfw") + self.assertEqual(model["family"], "yolo-det") + self.assertEqual(model["size"], tc_config.YOLO_DET_SIZE) + self.assertEqual(model["class_names"], ["person", "handjob"]) + + def test_build_cfg_det_defaults_and_overrides(self): + defaults = tc_cli.build_cfg({}) + self.assertEqual(defaults["det_conf"], tc_config.YOLO_DET_CONF_DEFAULT) + self.assertEqual(defaults["det_iou"], tc_config.YOLO_DET_IOU_DEFAULT) + cfg = tc_cli.build_cfg({"det_conf": "0.4", "det_iou": "0.6"}) + self.assertAlmostEqual(cfg["det_conf"], 0.4) + self.assertAlmostEqual(cfg["det_iou"], 0.6) + + + def test_classify_image_detects_act(self): + out = self._det_output([((320.0, 320.0, 100.0, 100.0), 8, 0.9)]) + model = { + "family": "yolo-det", + "size": tc_config.YOLO_DET_SIZE, + "session": _DetSession(out), + "input_name": "input", + "output_name": "output", + "class_names": tc_config.YOLO_DET_CLASS_NAMES, + } + cfg = { + "det_conf": tc_config.YOLO_DET_CONF_DEFAULT, + "det_iou": tc_config.YOLO_DET_IOU_DEFAULT, + } + img = Image.new("RGB", (1280, 720), (255, 255, 255)) + tags = tc_inference.classify_image(img, model, cfg) + self.assertEqual([name for name, _ in tags], ["blowjob"]) + + +class _DetSession: + """Minimal session stub returning a pre-built detection head.""" + + def __init__(self, output): + self._output = output + + def run(self, output_names, feeds): + return [self._output] + + +class TestSexPos(unittest.TestCase): + def setUp(self): + self.model = {"class_names": list(tc_config.SEXPOS_CLASS_NAMES)} + + def test_class_names(self): + self.assertEqual(len(tc_config.SEXPOS_CLASS_NAMES), 8) + self.assertEqual(tc_config.SEXPOS_CLASS_NAMES[0], "blowjob") + + def test_preprocess_shape_and_imagenet_normalisation(self): + img = Image.new("RGB", (400, 200), (255, 255, 255)) + arr = tc_preprocess.sexpos_preprocess(img) + self.assertEqual(arr.shape, (1, 3, 224, 224)) + self.assertAlmostEqual( + float(arr[0, 0, 0, 0]), (1.0 - 0.485) / 0.229, places=4 + ) + + def test_postprocess_picks_top_class(self): + logits = np.zeros(8, dtype=np.float32) + logits[7] = 5.0 + hits = tc_postprocess.sexpos_postprocess(logits, self.model) + self.assertEqual([name for name, _ in hits], ["solo"]) + self.assertGreater(hits[0][1], 0.9) + + def test_postprocess_respects_threshold(self): + logits = np.zeros(8, dtype=np.float32) + self.assertEqual( + tc_postprocess.sexpos_postprocess(logits, self.model, threshold=0.5), [] + ) + + def test_postprocess_margin(self): + logits = np.array([4.0, 3.5, 0, 0, 0, 0, 0, 0], dtype=np.float32) + self.assertEqual( + [name for name, _ in tc_postprocess.sexpos_postprocess(logits, self.model)], + ["blowjob"], + ) + self.assertEqual( + tc_postprocess.sexpos_postprocess(logits, self.model, margin=0.5), [] + ) + + def test_postprocess_name_fallback(self): + model = {"class_names": ["blowjob"]} + logits = np.zeros(8, dtype=np.float32) + logits[7] = 5.0 + self.assertEqual( + [name for name, _ in tc_postprocess.sexpos_postprocess(logits, model)], + ["class_7"], + ) + + def test_registry_loads_sexpos(self): + with tempfile.TemporaryDirectory() as tmp: + with open(os.path.join(tmp, "sex-position.onnx"), "w") as f: + f.write("x") + fake = _fake_onnxruntime(lambda path, providers=None: _FakeSession( + path, providers, input_shape=(1, 3, 224, 224), output_size=8)) + tc_models._SESSION_CACHE.clear() + self.addCleanup(tc_models._SESSION_CACHE.clear) + with mock.patch.object(tc_config, "MODEL_DIR", tmp), \ + mock.patch.dict(sys.modules, {"onnxruntime": fake}): + model = tc_models.get_model("sex-position") + self.assertEqual(model["family"], "sexpos") + self.assertEqual(model["size"], tc_config.SEXPOS_SIZE) + self.assertEqual(model["class_names"], tc_config.SEXPOS_CLASS_NAMES) + + def test_classify_image_uses_sexpos(self): + out = np.zeros((1, 8), dtype=np.float32) + out[0, 7] = 5.0 + model = { + "name": "sex-position", + "family": "sexpos", + "size": tc_config.SEXPOS_SIZE, + "session": _DetSession(out), + "input_name": "input", + "output_name": "output", + "class_names": tc_config.SEXPOS_CLASS_NAMES, + } + img = Image.new("RGB", (320, 240), (128, 128, 128)) + tags = tc_inference.classify_image(img, model, tc_cli.build_cfg({})) + self.assertEqual([name for name, _ in tags], ["solo"]) + + +class TestMultiModel(unittest.TestCase): + def test_parse_single_name(self): + self.assertEqual(tc_cli.parse_model_names("joytag", "d"), ["joytag"]) + + def test_parse_comma_semicolon_and_whitespace(self): + self.assertEqual( + tc_cli.parse_model_names("a, b;c d", "x"), ["a", "b", "c", "d"] + ) + + def test_parse_dedupes_and_keeps_order(self): + self.assertEqual(tc_cli.parse_model_names("a,b,a", "x"), ["a", "b"]) + + def test_parse_empty_uses_default(self): + self.assertEqual(tc_cli.parse_model_names("", "img"), ["img"]) + self.assertEqual(tc_cli.parse_model_names(None, "vid"), ["vid"]) + self.assertEqual(tc_cli.parse_model_names([], "d"), ["d"]) + + def test_parse_list_value(self): + self.assertEqual(tc_cli.parse_model_names(["a", "b", "a"], "d"), ["a", "b"]) + + def test_merge_tag_results_keeps_best_score(self): + merged = tc_inference.merge_tag_results([[("a", 0.4)], [("b", 0.9), ("a", 0.7)]]) + self.assertEqual(merged, [("b", 0.9), ("a", 0.7)]) + + def test_classify_with_models_runs_every_model(self): + img = Image.new("RGB", (8, 8), (0, 0, 0)) + models = [{"name": "m1"}, {"name": "m2"}] + calls = [] + + def fake(image, model, cfg): + calls.append(model["name"]) + if model["name"] == "m1": + return [("a", 0.4)] + return [("b", 0.9), ("a", 0.7)] + + with mock.patch.object(tc_inference, "classify_image", side_effect=fake): + tags = tc_inference.classify_with_models(img, models, {}) + self.assertEqual(calls, ["m1", "m2"]) + self.assertEqual(tags, [("b", 0.9), ("a", 0.7)]) + + def test_classify_with_models_single_returns_unchanged(self): + img = Image.new("RGB", (8, 8), (0, 0, 0)) + with mock.patch.object(tc_inference, "classify_image", return_value=[("x", 0.1), ("y", 0.9)] + ): + tags = tc_inference.classify_with_models(img, {"name": "m"}, {}) + self.assertEqual(tags, [("x", 0.1), ("y", 0.9)]) + + def test_classify_with_models_skips_failing_model(self): + img = Image.new("RGB", (8, 8), (0, 0, 0)) + models = [{"name": "bad"}, {"name": "good"}] + + def fake(image, model, cfg): + if model["name"] == "bad": + raise RuntimeError("boom") + return [("ok", 0.8)] + + with mock.patch.object(tc_inference, "classify_image", side_effect=fake), \ + mock.patch.object(tc_stashlog, "log", mock.MagicMock()): + tags = tc_inference.classify_with_models(img, models, {}) + self.assertEqual(tags, [("ok", 0.8)]) + + def test_infer_frames_with_model_list_merges(self): + with tempfile.TemporaryDirectory() as tmp: + paths = [] + for i in range(3): + path = os.path.join(tmp, f"frame_{i:03d}.jpg") + Image.new("RGB", (64, 64), (10, 20, 30)).save(path) + paths.append(path) + + models = [ + {"name": "sex-position", "family": "sexpos"}, + {"name": "joytag", "family": "joytag"}, + ] + + def fake(image, model, cfg): + if model["name"] == "sex-position": + return [("solo", 0.8)] + return [("1girl", 0.7), ("solo", 0.6)] + + with mock.patch.object(tc_inference, "classify_image", side_effect=fake): + tags = tc_inference.infer_frames(paths, models, tc_cli.build_cfg({})) + merged = dict(tags) + self.assertAlmostEqual(merged["solo"], 0.8, places=5) + self.assertAlmostEqual(merged["1girl"], 0.7, places=5) + + +class TestDebugLog(unittest.TestCase): + def tearDown(self): + tc_debug.debug_start(None, {}) + + def test_resolve_debug_path_disabled(self): + for value in (None, "", "false", "0", "no", "off", False): + self.assertIsNone(tc_debug.resolve_debug_path(value)) + + def test_resolve_debug_path_truthy_uses_default_file(self): + expected = os.path.join(tc_config.PLUGIN_DIR, tc_config.DEBUG_FILENAME) + for value in (True, "true", "1", "yes", "on"): + self.assertEqual(tc_debug.resolve_debug_path(value), expected) + + def test_resolve_debug_path_explicit(self): + self.assertEqual(tc_debug.resolve_debug_path("/tmp/x.json"), "/tmp/x.json") + self.assertEqual( + tc_debug.resolve_debug_path(" out/report.json "), "out/report.json" + ) + + def test_candidate_scores_vectors(self): + model = {"family": "yolo", "class_names": ["a", "b", "c"]} + cfg = {"debug_floor": 0.1, "debug_top": 2} + scores = np.array([0.05, 0.9, 0.3], dtype=np.float32) + + hits = tc_inference.candidate_scores("yolo", scores, model, cfg) + self.assertEqual([name for name, _score in hits], ["b", "c"]) + self.assertAlmostEqual(hits[0][1], 0.9, places=5) + self.assertAlmostEqual(hits[1][1], 0.3, places=5) + + # The floor can be raised for a single call. + raised = tc_inference.candidate_scores("yolo", scores, model, cfg, floor=0.5) + self.assertEqual([name for name, _score in raised], ["b"]) + + def test_candidate_scores_name_fallback(self): + model = {"family": "yolo"} + cfg = {"debug_floor": 0.0, "debug_top": 10} + scores = np.array([0.2, 0.4], dtype=np.float32) + + hits = tc_inference.candidate_scores("yolo", scores, model, cfg) + self.assertEqual([name for name, _score in hits], ["class_1", "class_0"]) + self.assertAlmostEqual(hits[0][1], 0.4, places=5) + self.assertAlmostEqual(hits[1][1], 0.2, places=5) + + def test_candidate_scores_empty(self): + self.assertEqual( + tc_inference.candidate_scores("yolo", np.array([], dtype=np.float32), {}, {}), + [], + ) + + def test_candidate_scores_detector_uses_the_floor(self): + calls = {} + + def fake_postprocess(scores, model, conf_threshold=None, + iou_threshold=None): + calls["conf"] = conf_threshold + calls["iou"] = iou_threshold + return [("person", 0.12), ("penis", 0.08)] + + model = {"family": "yolo-det", "class_names": ["person", "penis"]} + with mock.patch.object(tc_postprocess, "yolo_det_postprocess", side_effect=fake_postprocess + ): + hits = tc_inference.candidate_scores( + "yolo-det", np.zeros((1, 14, 10), dtype=np.float32), model, + {"debug_floor": 0.05, "debug_top": 1, "det_iou": 0.5}, + ) + + self.assertEqual(calls, {"conf": 0.05, "iou": 0.5}) + self.assertEqual(hits, [("person", 0.12)]) + + def test_aggregate_candidates_keeps_max_per_tag(self): + merged = tc_debug.aggregate_candidates([ + [("a", 0.4), ("b", 0.9)], + [("a", 0.7), ("c", 0.2)], + ]) + self.assertEqual(merged, [("b", 0.9), ("a", 0.7), ("c", 0.2)]) + self.assertEqual(tc_debug.aggregate_candidates([]), []) + + def test_finish_debug_entry_marks_accepted_candidates(self): + entry = tc_debug._debug_entry( + {"name": "m", "family": "yolo"}, {"yolo_threshold": 0.5} + ) + entry["per_frame"] = [[("a", 0.9), ("b", 0.3)]] + tc_debug._finish_debug_entry(entry, [("a", 0.9)]) + + self.assertEqual(entry["tags"], ["a"]) + self.assertEqual( + entry["aggregate"][0], + {"tag": "a", "score": 0.9, "accepted": True}, + ) + self.assertEqual(entry["aggregate"][1]["tag"], "b") + self.assertFalse(entry["aggregate"][1]["accepted"]) + + def test_debug_report_is_written(self): + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "sub", "report.json") + cfg = {"debug_log": path, "debug_floor": 0.05, "hf_token": "secret"} + self.assertTrue( + tc_debug.debug_start(tc_debug.resolve_debug_path(cfg["debug_log"]), cfg) + ) + tc_debug.debug_add_item( + {"id": 1, "models": [], "applied_tags": ["AI: x"]} + ) + with open(path, encoding="utf-8") as handle: + data = json.load(handle) + + self.assertEqual(data["plugin"], "TagAI") + self.assertEqual(data["format"], 1) + self.assertEqual(len(data["items"]), 1) + self.assertEqual(data["items"][0]["id"], 1) + self.assertEqual(data["items"][0]["applied_tags"], ["AI: x"]) + self.assertNotIn("hf_token", data["settings"]) + + def test_debug_add_item_without_path_is_noop(self): + tc_debug.debug_start(None, {}) + tc_debug.debug_add_item({"id": 1}) + self.assertEqual(tc_debug._DEBUG["items"], []) + + def test_infer_frames_collects_debug_records(self): + with tempfile.TemporaryDirectory() as tmp: + paths = [] + for i in range(2): + path = os.path.join(tmp, f"frame_{i}.jpg") + Image.new("RGB", (64, 64), (1, 2, 3)).save(path) + paths.append(path) + + session = mock.MagicMock() + session.run.return_value = [ + np.array([[0.8, 0.1], [0.8, 0.1]], dtype=np.float32) + ] + model = { + "name": "yolo26n-cls-porn", "family": "yolo", "size": 224, + "session": session, "input_name": "input", + "output_name": "output", + "class_names": ["porn", "safe"], "thresholds": None, + } + cfg = tc_cli.build_cfg({"debug_log": "report.json", "yolo_margin": "0.0"}) + debug = {} + + tags = tc_inference.infer_frames(paths, model, cfg, debug=debug) + entry = debug["yolo26n-cls-porn"] + + self.assertEqual(entry["family"], "yolo") + self.assertEqual(len(entry["per_frame"]), 2) + self.assertEqual(entry["per_frame"][0][0][0], "porn") + self.assertEqual(entry["tags"], [name for name, _score in tags]) + self.assertEqual(entry["tags"], ["porn"]) + self.assertTrue(entry["aggregate"][0]["accepted"]) + + def test_classify_with_models_collects_debug_records(self): + img = Image.new("RGB", (8, 8), (0, 0, 0)) + models = [{"name": "m1", "family": "yolo", "class_names": ["a", "b"]}] + cfg = {"debug_floor": 0.0, "debug_top": 5} + debug = [] + + with mock.patch.object(tc_inference, "run_inference", + return_value=("yolo", np.array([0.9, 0.2], dtype=np.float32)), + ), mock.patch.object(tc_inference, "postprocess_scores", return_value=[("a", 0.9)] + ): + tags = tc_inference.classify_with_models(img, models, cfg, debug=debug) + + self.assertEqual(tags, [("a", 0.9)]) + self.assertEqual(len(debug), 1) + self.assertEqual(debug[0]["name"], "m1") + self.assertEqual(debug[0]["tags"], ["a"]) + per_frame = debug[0]["per_frame"] + self.assertEqual([name for name, _score in per_frame[0]], ["a", "b"]) + self.assertAlmostEqual(per_frame[0][0][1], 0.9, places=5) + + +class TestSceneMarkers(unittest.TestCase): + def test_parse_marker_mode(self): + for value in ("", None, "off", "OFF", "false", "0", "no", "none", + "disabled"): + self.assertEqual(tc_markers.parse_marker_mode(value), "off") + for value in ("ranges", "Range", "runs"): + self.assertEqual(tc_markers.parse_marker_mode(value), "ranges") + for value in ("best", "TOP", "peak"): + self.assertEqual(tc_markers.parse_marker_mode(value), "best") + with mock.patch.object(tc_stashlog, "log", mock.MagicMock()): + self.assertEqual(tc_markers.parse_marker_mode("nonsense"), "off") + + def test_frame_spacing(self): + self.assertEqual(tc_markers._frame_spacing([]), 0.0) + self.assertEqual(tc_markers._frame_spacing([5.0]), 0.0) + self.assertAlmostEqual(tc_markers._frame_spacing([0.0, 10.0, 20.0]), 10.0) + # Median of the gaps, without interpolation. + self.assertAlmostEqual(tc_markers._frame_spacing([0.0, 1.0, 5.0, 6.0]), 1.0) + + def test_combine_hits_merges_models_and_keeps_the_best(self): + hits_out = { + "m1": [[("a", 0.4)], []], + "m2": [[("a", 0.9), ("b", 0.6)], [("c", 0.3)]], + } + combined = tc_markers.combine_hits(hits_out, 3) + self.assertEqual(len(combined), 3) + self.assertAlmostEqual(combined[0]["a"], 0.9) + self.assertAlmostEqual(combined[0]["b"], 0.6) + self.assertAlmostEqual(combined[1]["c"], 0.3) + self.assertEqual(combined[2], {}) + # Rows past the frame count are ignored. + self.assertEqual(len(tc_markers.combine_hits(hits_out, 1)), 1) + self.assertEqual(tc_markers.combine_hits(None, 2), [{}, {}]) + + def test_marker_span_pads_the_ends(self): + span = tc_markers._marker_span("a", 10.0, 12.0, 0.8, ["a", "b"], + {"b": 0.5}, 2.0, None) + self.assertEqual(span, {"tag": "a", "start": 8.0, "end": 14.0, + "score": 0.8, "tags": ["b"]}) + + def test_marker_span_minimum_length_and_clamping(self): + span = tc_markers._marker_span("a", 10.0, 10.0, 0.9, ["a"], {}, 0.0, None) + self.assertAlmostEqual(span["end"] - span["start"], + tc_config.MARKER_MIN_LENGTH) + + # Clamped so the marker never runs past the end of the video. + span = tc_markers._marker_span("a", 9.5, 9.6, 0.9, ["a"], {}, 0.0, 10.0) + self.assertEqual(span["start"], 9.0) + self.assertEqual(span["end"], 10.0) + + # Never starts before the beginning either. + span = tc_markers._marker_span("a", 1.0, 1.0, 0.9, ["a"], {}, 30.0, 100.0) + self.assertEqual(span["start"], 0.0) + + def test_build_markers_ranges_merges_adjacent_frames(self): + markers = tc_markers.build_markers( + [("a", 0.9)], + [{"a": 0.6}, {"a": 0.7}, {}, {"a": 0.9}, {"a": 0.5}], + [0.0, 10.0, 20.0, 30.0, 40.0], + 100.0, + pad=0.0, + ) + self.assertEqual( + [(m["tag"], m["start"], m["end"], m["score"]) for m in markers], + [("a", 0.0, 10.0, 0.7), ("a", 30.0, 40.0, 0.9)], + ) + + def test_build_markers_merges_when_padded_spans_would_overlap(self): + tags = [("a", 0.9)] + hits = [{"a": 0.9}, {}, {"a": 0.8}] + times = [0.0, 10.0, 20.0] + + merged = tc_markers.build_markers(tags, hits, times, 100.0, pad=10.0) + self.assertEqual(len(merged), 1) + self.assertEqual((merged[0]["start"], merged[0]["end"]), (0.0, 30.0)) + + split = tc_markers.build_markers(tags, hits, times, 100.0, pad=5.0) + self.assertEqual(len(split), 2) + + def test_build_markers_best_mode_uses_the_peak_frame(self): + markers = tc_markers.build_markers( + [("a", 0.9)], + [{"a": 0.7}, {"a": 0.9}, {"a": 0.6}], + [0.0, 10.0, 20.0], + 100.0, + mode="best", + pad=0.0, + ) + self.assertEqual(len(markers), 1) + self.assertEqual(markers[0]["start"], 10.0) + self.assertEqual(markers[0]["end"], 11.0) + self.assertAlmostEqual(markers[0]["score"], 0.9) + + def test_build_markers_ignores_unapplied_tags_and_collects_co_tags(self): + markers = tc_markers.build_markers( + [("a", 0.9), ("b", 0.6)], + [{"a": 0.9, "b": 0.6, "c": 0.99}], + [5.0], + 100.0, + pad=0.0, + ) + self.assertEqual([m["tag"] for m in markers], ["a", "b"]) + # "c" was never applied to the scene, so it is not a co-tag either. + self.assertEqual(markers[0]["tags"], ["b"]) + self.assertEqual(markers[1]["tags"], ["a"]) + + def test_build_markers_without_hits(self): + self.assertEqual(tc_markers.build_markers([], [], [], 10.0), []) + self.assertEqual( + tc_markers.build_markers([("a", 0.9)], [], [1.0], 10.0), []) + + def test_get_scene_markers_parses_the_response(self): + payload = {"data": {"findScene": {"scene_markers": [ + {"seconds": 12.5, "primary_tag": {"id": "7"}}, + {"seconds": 0, "primary_tag": None}, + ]}}} + with mock.patch.object(tc_stash_client, "gql", return_value=payload): + markers = tc_markers.get_scene_markers("3") + self.assertEqual(markers, [ + {"primary_tag_id": "7", "seconds": 12.5}, + {"primary_tag_id": None, "seconds": 0.0}, + ]) + + def test_create_scene_marker_payload(self): + calls = [] + + def fake_gql(query, variables=None): + calls.append(variables) + return {"data": {"sceneMarkerCreate": {"id": "11"}}} + + with mock.patch.object(tc_stash_client, "gql", fake_gql): + ok = tc_markers.create_scene_marker( + "2", "5", 10.1234, 20.5678, "t", ["6"]) + + self.assertTrue(ok) + self.assertEqual(calls[0]["input"], { + "scene_id": "2", "primary_tag_id": "5", "seconds": 10.123, + "end_seconds": 20.568, "title": "t", "tag_ids": ["6"], + }) + + # Without co-tags the key is omitted and a failed create returns False. + with mock.patch.object(tc_stash_client, "gql", + return_value={"data": {}}) as fake: + self.assertFalse( + tc_markers.create_scene_marker("2", "5", 0.0, 1.0, "t")) + self.assertNotIn("tag_ids", fake.call_args.args[1]["input"]) + + def test_create_scene_markers_off_by_default(self): + with mock.patch.object(tc_stash_client, "gql", mock.MagicMock()) as fake: + created = tc_markers.create_scene_markers( + {"id": "1"}, [("a", 0.9)], [{"a": 0.9}], [10.0], 100.0, + {"markers": "off"}, "AI: ") + self.assertEqual(created, []) + fake.assert_not_called() + + def test_create_scene_markers_creates_and_dedupes(self): + created_calls = [] + + def fake_create(scene_id, tag_id, start, end, title, tag_ids=None): + created_calls.append((scene_id, tag_id, start, end, title, tag_ids)) + return True + + cfg = {"markers": "ranges", "marker_pad": 0.0, "marker_max": 50} + scene = {"id": "9"} + with mock.patch.object(tc_stash_client, "ensure_tag", + side_effect=lambda name: "id-" + name), \ + mock.patch.object(tc_markers, "get_scene_markers", + return_value=[]), \ + mock.patch.object(tc_markers, "create_scene_marker", fake_create): + created = tc_markers.create_scene_markers( + scene, [("a", 0.9)], [{"a": 0.9}, {"a": 0.8}], [0.0, 10.0], + 100.0, cfg, "AI: ") + + self.assertEqual(len(created), 1) + self.assertEqual(created[0]["start"], 0.0) + self.assertEqual(created_calls, + [("9", "id-AI: a", 0.0, 10.0, "a", [])]) + + # A re-run sees the marker that already exists and skips it. + with mock.patch.object(tc_stash_client, "ensure_tag", + side_effect=lambda name: "id-" + name), \ + mock.patch.object(tc_markers, "get_scene_markers", + return_value=[{"primary_tag_id": "id-AI: a", + "seconds": 0.5}]), \ + mock.patch.object(tc_markers, "create_scene_marker", fake_create): + created = tc_markers.create_scene_markers( + scene, [("a", 0.9)], [{"a": 0.9}], [0.0], 100.0, cfg, "AI: ") + self.assertEqual(created, []) + + def test_create_scene_markers_honours_marker_max(self): + created = [] + + def fake_create(scene_id, tag_id, start, end, title, tag_ids=None): + created.append(title) + return True + + cfg = {"markers": "best", "marker_pad": 0.0, "marker_max": 1} + with mock.patch.object(tc_stash_client, "ensure_tag", + side_effect=lambda name: "id-" + name), \ + mock.patch.object(tc_markers, "get_scene_markers", + return_value=[]), \ + mock.patch.object(tc_markers, "create_scene_marker", fake_create), \ + mock.patch.object(tc_stashlog, "log", mock.MagicMock()): + spans = tc_markers.create_scene_markers( + {"id": "4"}, [("a", 0.9), ("b", 0.8)], + [{"a": 0.9}, {"b": 0.8}], [0.0, 10.0], 100.0, cfg, "AI: ") + + self.assertEqual(len(spans), 1) + self.assertEqual(created, ["a"]) + + def test_create_scene_markers_default_pad_is_the_frame_spacing(self): + created = [] + + def fake_create(scene_id, tag_id, start, end, title, tag_ids=None): + created.append((start, end)) + return True + + cfg = {"markers": "ranges", "marker_pad": None, "marker_max": 10} + with mock.patch.object(tc_stash_client, "ensure_tag", + side_effect=lambda name: "id-" + name), \ + mock.patch.object(tc_markers, "get_scene_markers", + return_value=[]), \ + mock.patch.object(tc_markers, "create_scene_marker", fake_create): + spans = tc_markers.create_scene_markers( + {"id": "5"}, [("a", 0.9)], [{"a": 0.9}, {}, {"a": 0.8}], + [0.0, 10.0, 20.0], 100.0, cfg, "AI: ") + + self.assertEqual(len(spans), 1) + self.assertEqual(created, [(0.0, 30.0)]) + + + def test_parse_tag_filter(self): + self.assertEqual(tc_tags.parse_tag_filter(""), []) + self.assertEqual(tc_tags.parse_tag_filter(None), []) + self.assertEqual(tc_tags.parse_tag_filter("a, b ; c"), ["a", "b", "c"]) + # The prefix is optional and matching is case-insensitive. + self.assertEqual( + tc_tags.parse_tag_filter("AI: Cowgirl Position, sex from behind"), + ["cowgirl position", "sex from behind"], + ) + # Whitespace never separates a multi-word tag name. + self.assertEqual( + tc_tags.parse_tag_filter("cowgirl position"), ["cowgirl position"] + ) + self.assertEqual(tc_tags.parse_tag_filter(["B", "a", "b"]), ["a", "b"]) + + def test_filter_tags_filters_scene_tags(self): + tags = [("a", 0.9), ("b", 0.8), ("c", 0.7)] + self.assertEqual(tc_tags.filter_tags(tags), tags) + self.assertEqual( + tc_tags.filter_tags(tags, ["a", "b"]), [("a", 0.9), ("b", 0.8)] + ) + self.assertEqual( + tc_tags.filter_tags(tags, None, ["b"]), [("a", 0.9), ("c", 0.7)] + ) + self.assertEqual(tc_tags.filter_tags(tags, ["a", "b"], ["a"]), [("b", 0.8)]) + + def test_filter_marker_tags(self): + tags = [("a", 0.9), ("b", 0.8), ("c", 0.7)] + + self.assertEqual(tc_tags.filter_marker_tags(tags), tags) + self.assertEqual( + tc_tags.filter_marker_tags(tags, None, ["b"]), + [("a", 0.9), ("c", 0.7)], + ) + self.assertEqual( + tc_tags.filter_marker_tags(tags, ["a", "b"]), + [("a", 0.9), ("b", 0.8)], + ) + # The blacklist wins over the whitelist. + self.assertEqual( + tc_tags.filter_marker_tags(tags, ["a", "b"], ["a"]), [("b", 0.8)] + ) + # Matching is case-insensitive. + self.assertEqual( + tc_tags.filter_marker_tags( + [("Cowgirl Position", 0.9)], None, ["cowgirl position"] + ), + [], + ) + + def test_create_scene_markers_blacklist_skips_the_tag(self): + created_calls = [] + + def fake_create(scene_id, tag_id, start, end, title, tag_ids=None): + created_calls.append((scene_id, tag_id, start, end, title, tag_ids)) + return True + + cfg = { + "markers": "ranges", "marker_pad": 0.0, "marker_max": 50, + "marker_whitelist": [], "marker_blacklist": ["noise"], + } + with mock.patch.object(tc_stash_client, "ensure_tag", + side_effect=lambda name: "id-" + name), \ + mock.patch.object(tc_markers, "get_scene_markers", + return_value=[]), \ + mock.patch.object(tc_markers, "create_scene_marker", + side_effect=fake_create): + created = tc_markers.create_scene_markers( + {"id": "7"}, [("a", 0.9), ("noise", 0.95)], + [{"a": 0.9, "noise": 0.95}], [0.0], 100.0, cfg, "AI: ", + ) + + self.assertEqual(len(created), 1) + self.assertEqual(created[0]["tag"], "a") + # "noise" creates no marker and is not attached as a co-tag either. + self.assertEqual( + created_calls, [("7", "id-AI: a", 0.0, 1.0, "a", [])] + ) + + def test_create_scene_markers_whitelist_restricts_the_tags(self): + created_calls = [] + + def fake_create(scene_id, tag_id, start, end, title, tag_ids=None): + created_calls.append((scene_id, tag_id, start, end, title, tag_ids)) + return True + + cfg = { + "markers": "ranges", "marker_pad": 0.0, "marker_max": 50, + "marker_whitelist": ["a"], "marker_blacklist": [], + } + with mock.patch.object(tc_stash_client, "ensure_tag", + side_effect=lambda name: "id-" + name), \ + mock.patch.object(tc_markers, "get_scene_markers", + return_value=[]), \ + mock.patch.object(tc_markers, "create_scene_marker", + side_effect=fake_create): + created = tc_markers.create_scene_markers( + {"id": "8"}, [("a", 0.9), ("b", 0.8)], + [{"a": 0.9, "b": 0.8}], [0.0], 100.0, cfg, "AI: ", + ) + + self.assertEqual([marker["tag"] for marker in created], ["a"]) + # "b" is not attached to the "a" marker either. + self.assertEqual(created_calls, [("8", "id-AI: a", 0.0, 1.0, "a", [])]) + + def test_create_scene_markers_whitelist_can_leave_nothing(self): + cfg = { + "markers": "ranges", "marker_pad": 0.0, "marker_max": 50, + "marker_whitelist": ["z"], "marker_blacklist": [], + } + with mock.patch.object(tc_stash_client, "gql") as gql: + created = tc_markers.create_scene_markers( + {"id": "8"}, [("a", 0.9)], [{"a": 0.9}], [0.0], 100.0, + cfg, "AI: ", + ) + + self.assertEqual(created, []) + gql.assert_not_called() + + + def test_build_markers_merges_hits_within_the_gap(self): + tags = [("a", 0.9)] + hits = [{"a": 0.9}, {}, {"a": 0.8}] + times = [0.0, 10.0, 30.0] + + apart = tc_markers.build_markers(tags, hits, times, 100.0, pad=0.0) + self.assertEqual(len(apart), 2) + self.assertEqual( + [(marker["start"], marker["end"]) for marker in apart], + [(0.0, 1.0), (30.0, 31.0)], + ) + + merged = tc_markers.build_markers(tags, hits, times, 100.0, pad=0.0, gap=60.0) + self.assertEqual(len(merged), 1) + self.assertEqual(merged[0]["start"], 0.0) + self.assertEqual(merged[0]["end"], 30.0) + self.assertAlmostEqual(merged[0]["score"], 0.9, places=5) + + def test_build_markers_gap_never_merges_different_tags(self): + markers = tc_markers.build_markers( + [("a", 0.9), ("b", 0.9)], [{"a": 0.9}, {"b": 0.9}], [0.0, 30.0], + 100.0, pad=0.0, gap=60.0, + ) + self.assertEqual([marker["tag"] for marker in markers], ["a", "b"]) + self.assertEqual(len(markers), 2) + + +class TestFrameHits(unittest.TestCase): + def test_failed_batch_pads_the_rows(self): + with tempfile.TemporaryDirectory() as tmp: + paths = [] + for i in range(2): + path = os.path.join(tmp, f"frame_{i}.jpg") + Image.new("RGB", (64, 64), (1, 2, 3)).save(path) + paths.append(path) + + session = mock.MagicMock() + session.run.side_effect = [ + RuntimeError("boom"), + [np.array([[0.9, 0.1]], dtype=np.float32)], + ] + model = { + "name": "cls", "family": "yolo", "size": 224, + "session": session, "input_name": "input", + "output_name": "output", + "class_names": ["porn", "safe"], "thresholds": None, + } + cfg = tc_cli.build_cfg({"yolo_margin": "0.0", + "infer_batch_size": "1"}) + hits = {} + with mock.patch.object(tc_stashlog, "log", mock.MagicMock()): + tags = tc_inference.infer_frames(paths, model, cfg, hits_out=hits) + + self.assertEqual(len(hits["cls"]), 2) + self.assertEqual(hits["cls"][0], []) + self.assertEqual(hits["cls"][1][0][0], "porn") + self.assertIn("porn", [name for name, _score in tags]) + + def test_failed_frame_keeps_the_rows_aligned(self): + with tempfile.TemporaryDirectory() as tmp: + paths = [] + for i in range(3): + path = os.path.join(tmp, f"frame_{i}.jpg") + Image.new("RGB", (64, 64), (1, 2, 3)).save(path) + paths.append(path) + + calls = [] + + def fake(image, model, cfg): + calls.append(1) + if len(calls) == 2: + raise RuntimeError("boom") + return [("solo", 0.9)] + + hits = {} + with mock.patch.object(tc_inference, "classify_image", side_effect=fake), \ + mock.patch.object(tc_stashlog, "log", mock.MagicMock()): + tags = tc_inference.infer_frames( + paths, {"name": "m", "family": "sexpos"}, + tc_cli.build_cfg({}), hits_out=hits) + + self.assertEqual([len(row) for row in hits["m"]], [1, 0, 1]) + self.assertEqual([name for name, _score in tags], ["solo"]) + self.assertAlmostEqual(tags[0][1], 0.9, places=5) + + def test_hits_out_is_optional(self): + with mock.patch.object(tc_inference, "classify_image", + return_value=[("solo", 0.9)]): + tags = tc_inference.infer_frames( + ["/no/such/frame.jpg"], {"name": "m", "family": "sexpos"}, + tc_cli.build_cfg({})) + + self.assertEqual(tags, []) + + +class TestBuildCfg(unittest.TestCase): + def test_defaults(self): + cfg = tc_cli.build_cfg({}) + self.assertEqual(cfg["frame_count"], tc_config.FRAME_COUNT_DEFAULT) + self.assertEqual( + cfg["sexpos_threshold"], tc_config.SEXPOS_THRESHOLD_DEFAULT + ) + self.assertEqual(cfg["sexpos_margin"], tc_config.SEXPOS_MARGIN_DEFAULT) + self.assertEqual(cfg["det_conf"], tc_config.YOLO_DET_CONF_DEFAULT) + self.assertEqual(cfg["det_iou"], tc_config.YOLO_DET_IOU_DEFAULT) + self.assertEqual(cfg["frame_interval"], tc_config.FRAME_INTERVAL_DEFAULT) + self.assertEqual(cfg["max_frames"], tc_config.MAX_FRAMES_DEFAULT) + self.assertEqual(cfg["frame_workers"], tc_config.FRAME_WORKERS_DEFAULT) + self.assertEqual(cfg["infer_batch_size"], tc_config.INFER_BATCH_SIZE_DEFAULT) + self.assertEqual(cfg["ffmpeg_path"], "ffmpeg") + self.assertFalse(cfg["tag_characters"]) + self.assertTrue(cfg["yolo_multi"]) + self.assertEqual(cfg["yolo_min_votes"], tc_config.YOLO_MIN_VOTES_DEFAULT) + self.assertEqual(cfg["yolo_share"], tc_config.YOLO_SHARE_DEFAULT) + self.assertEqual(cfg["yolo_top_k"], tc_config.YOLO_TOP_K_DEFAULT) + self.assertEqual(cfg["yolo_top_k"], 1) + + def test_overrides_and_bool(self): + cfg = tc_cli.build_cfg({ + "frame_count": "2", "tag_threshold": "0.7", + "tag_characters": "true", "ffmpeg_path": "/opt/ffmpeg", + }) + self.assertEqual(cfg["frame_count"], 2) + self.assertAlmostEqual(cfg["tag_threshold"], 0.7) + self.assertTrue(cfg["tag_characters"]) + self.assertEqual(cfg["ffmpeg_path"], "/opt/ffmpeg") + + def test_sampling_and_voting_overrides(self): + cfg = tc_cli.build_cfg({ + "frame_interval": "3.5", "max_frames": "12", + "yolo_multi": "false", "yolo_min_votes": "4", "yolo_share": "0.6", + "frame_workers": "2", "infer_batch_size": "16", "yolo_top_k": "2", + }) + self.assertAlmostEqual(cfg["frame_interval"], 3.5) + self.assertEqual(cfg["max_frames"], 12) + self.assertFalse(cfg["yolo_multi"]) + self.assertEqual(cfg["yolo_min_votes"], 4) + self.assertAlmostEqual(cfg["yolo_share"], 0.6) + self.assertEqual(cfg["frame_workers"], 2) + self.assertEqual(cfg["infer_batch_size"], 16) + self.assertEqual(cfg["yolo_top_k"], 2) + + def test_download_defaults_and_overrides(self): + cfg = tc_cli.build_cfg({}) + self.assertTrue(cfg["auto_download"]) + self.assertIsNone(cfg["hf_token"]) + + cfg = tc_cli.build_cfg({"auto_download": "false", "hf_token": "tok"}) + self.assertFalse(cfg["auto_download"]) + self.assertEqual(cfg["hf_token"], "tok") + + def test_debug_defaults_and_overrides(self): + cfg = tc_cli.build_cfg({}) + self.assertIsNone(cfg["debug_log"]) + self.assertEqual(cfg["debug_floor"], tc_config.DEBUG_FLOOR_DEFAULT) + self.assertEqual(cfg["debug_top"], tc_config.DEBUG_TOP_DEFAULT) + + cfg = tc_cli.build_cfg({ + "debug_log": "report.json", "debug_floor": "0.2", "debug_top": "3", + }) + self.assertEqual(cfg["debug_log"], "report.json") + self.assertAlmostEqual(cfg["debug_floor"], 0.2) + self.assertEqual(cfg["debug_top"], 3) + + def test_marker_defaults_and_overrides(self): + cfg = tc_cli.build_cfg({}) + self.assertEqual(cfg["markers"], "off") + self.assertIsNone(cfg["marker_pad"]) + self.assertEqual(cfg["marker_max"], tc_config.MARKER_MAX_DEFAULT) + + cfg = tc_cli.build_cfg({ + "markers": "best", "marker_pad": "2.5", "marker_max": "3", + }) + self.assertEqual(cfg["markers"], "best") + self.assertAlmostEqual(cfg["marker_pad"], 2.5) + self.assertEqual(cfg["marker_max"], 3) + + def test_marker_filter_defaults_and_overrides(self): + cfg = tc_cli.build_cfg({}) + self.assertEqual(cfg["marker_whitelist"], []) + self.assertEqual(cfg["marker_blacklist"], []) + + cfg = tc_cli.build_cfg({ + "marker_whitelist": "AI: Cowgirl Position,sex from behind", + "marker_blacklist": "blurry, 1girl", + }) + self.assertEqual( + cfg["marker_whitelist"], ["cowgirl position", "sex from behind"] + ) + self.assertEqual(cfg["marker_blacklist"], ["1girl", "blurry"]) + # The cfg is echoed into the debug report, so it has to be serialisable. + json.dumps(cfg) + + + def test_marker_gap_defaults_and_overrides(self): + self.assertEqual(tc_cli.build_cfg({})["marker_gap"], 0.0) + self.assertAlmostEqual( + tc_cli.build_cfg({"marker_gap": "45"})["marker_gap"], 45.0 + ) + + + def test_tag_filter_defaults_and_overrides(self): + cfg = tc_cli.build_cfg({}) + self.assertEqual(cfg["tag_whitelist"], []) + self.assertEqual(cfg["tag_blacklist"], []) + + cfg = tc_cli.build_cfg({ + "tag_whitelist": "AI: Cowgirl Position,sex from behind", + "tag_blacklist": "blurry, 1girl", + }) + self.assertEqual(cfg["tag_whitelist"], + ["cowgirl position", "sex from behind"]) + self.assertEqual(cfg["tag_blacklist"], ["1girl", "blurry"]) + + +if __name__ == "__main__": + unittest.main()