tagger support

app.py

@@ -9,6 +9,7 @@ from huggingface_hub import HfApi, snapshot_download
 from pnginfo import read_info_from_image, send_paras
 from images_history import img_history_ui
 from director_tools import director_ui, send_outputs
+from tagger import tagger_ui
 from utils import set_token, generate_novelai_image, image_from_bytes, get_remain_anlas, calculate_cost
 
 client_config = toml.load("config.toml")['client']
@@ -258,6 +259,8 @@ def ui():
                 from_t2i, send_i2i, send_inp, send_vib, in_image, out_image, d_index = director_ui()
             with gr.TabItem("图片信息读取"):
                 png2main, png_items, info, read_image = util_ui()
+            with gr.TabItem("Tagger反推"):
+                tags, tagger2main = tagger_ui()
             with gr.TabItem("云端图片浏览") as tab:
                 gallery, h_index, gal2main, gal_items, history2ref, history2i2i, history2inp, history2dtl = img_history_ui(tab)
             with gr.TabItem("设置"):
@@ -278,6 +281,7 @@ def ui():
             send_jump_select(history2inp, send_outputs, [gallery, h_index], paras[22], "client_ui_main", others[6], gal_items, "inp_block")
             send_jump_select(history2ref, (lambda l, i: None if i == -1 else [l[i]]), [gallery, h_index], paras[15], "client_ui_main", others[5], gal_items)
             send_and_jump(history2dtl, send_outputs, [gallery, h_index], in_image, "client_ui_dtool", gal_items)
+            send_and_jump(tagger2main, (lambda x: x), tags, paras[0], "client_ui_main", tags)
             read_image.change(read_info_from_image, inputs=read_image, outputs=[info, png_items])
     return website
 

requirements.txt

@@ -4,4 +4,7 @@ pillow
 numpy
 gradio==4.38.1
 toml
-piexif
+piexif
+pillow>=9.0.0
+onnxruntime>=1.12.0
+huggingface-hub

tagger.py

@@ -0,0 +1,297 @@
+import argparse
+
+import gradio as gr
+import huggingface_hub
+import numpy as np
+import onnxruntime as rt
+import pandas as pd
+from PIL import Image
+
+# Dataset v3 series of models:
+SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
+CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
+VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
+VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
+EVA02_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
+
+# Dataset v2 series of models:
+MOAT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
+SWIN_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
+CONV_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
+CONV2_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
+VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
+
+# Files to download from the repos
+MODEL_FILENAME = "model.onnx"
+LABEL_FILENAME = "selected_tags.csv"
+
+# https://github.com/toriato/stable-diffusion-webui-wd14-tagger/blob/a9eacb1eff904552d3012babfa28b57e1d3e295c/tagger/ui.py#L368
+kaomojis = [
+    "0_0",
+    "(o)_(o)",
+    "+_+",
+    "+_-",
+    "._.",
+    "<o>_<o>",
+    "<|>_<|>",
+    "=_=",
+    ">_<",
+    "3_3",
+    "6_9",
+    ">_o",
+    "@_@",
+    "^_^",
+    "o_o",
+    "u_u",
+    "x_x",
+    "|_|",
+    "||_||",
+]
+
+def load_labels(dataframe) -> list[str]:
+    name_series = dataframe["name"]
+    name_series = name_series.map(
+        lambda x: x.replace("_", " ") if x not in kaomojis else x
+    )
+    tag_names = name_series.tolist()
+
+    rating_indexes = list(np.where(dataframe["category"] == 9)[0])
+    general_indexes = list(np.where(dataframe["category"] == 0)[0])
+    character_indexes = list(np.where(dataframe["category"] == 4)[0])
+    return tag_names, rating_indexes, general_indexes, character_indexes
+
+
+def mcut_threshold(probs):
+    """
+    Maximum Cut Thresholding (MCut)
+    Largeron, C., Moulin, C., & Gery, M. (2012). MCut: A Thresholding Strategy
+     for Multi-label Classification. In 11th International Symposium, IDA 2012
+     (pp. 172-183).
+    """
+    sorted_probs = probs[probs.argsort()[::-1]]
+    difs = sorted_probs[:-1] - sorted_probs[1:]
+    t = difs.argmax()
+    thresh = (sorted_probs[t] + sorted_probs[t + 1]) / 2
+    return thresh
+
+
+class Predictor:
+    def __init__(self):
+        self.model_target_size = None
+        self.last_loaded_repo = None
+
+    def download_model(self, model_repo):
+        csv_path = huggingface_hub.hf_hub_download(
+            model_repo,
+            LABEL_FILENAME,
+        )
+        model_path = huggingface_hub.hf_hub_download(
+            model_repo,
+            MODEL_FILENAME,
+        )
+        return csv_path, model_path
+
+    def load_model(self, model_repo):
+        if model_repo == self.last_loaded_repo:
+            return
+
+        csv_path, model_path = self.download_model(model_repo)
+
+        tags_df = pd.read_csv(csv_path)
+        sep_tags = load_labels(tags_df)
+
+        self.tag_names = sep_tags[0]
+        self.rating_indexes = sep_tags[1]
+        self.general_indexes = sep_tags[2]
+        self.character_indexes = sep_tags[3]
+
+        model = rt.InferenceSession(model_path)
+        _, height, width, _ = model.get_inputs()[0].shape
+        self.model_target_size = height
+
+        self.last_loaded_repo = model_repo
+        self.model = model
+
+    def prepare_image(self, image):
+        target_size = self.model_target_size
+
+        canvas = Image.new("RGBA", image.size, (255, 255, 255))
+        canvas.alpha_composite(image)
+        image = canvas.convert("RGB")
+
+        # Pad image to square
+        image_shape = image.size
+        max_dim = max(image_shape)
+        pad_left = (max_dim - image_shape[0]) // 2
+        pad_top = (max_dim - image_shape[1]) // 2
+
+        padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
+        padded_image.paste(image, (pad_left, pad_top))
+
+        # Resize
+        if max_dim != target_size:
+            padded_image = padded_image.resize(
+                (target_size, target_size),
+                Image.BICUBIC,
+            )
+
+        # Convert to numpy array
+        image_array = np.asarray(padded_image, dtype=np.float32)
+
+        # Convert PIL-native RGB to BGR
+        image_array = image_array[:, :, ::-1]
+
+        return np.expand_dims(image_array, axis=0)
+
+    def predict(
+        self,
+        image,
+        model_repo,
+        general_thresh,
+        general_mcut_enabled,
+        character_thresh,
+        character_mcut_enabled,
+    ):
+        self.load_model(model_repo)
+
+        image = self.prepare_image(image)
+
+        input_name = self.model.get_inputs()[0].name
+        label_name = self.model.get_outputs()[0].name
+        preds = self.model.run([label_name], {input_name: image})[0]
+
+        labels = list(zip(self.tag_names, preds[0].astype(float)))
+
+        # First 4 labels are actually ratings: pick one with argmax
+        ratings_names = [labels[i] for i in self.rating_indexes]
+        rating = dict(ratings_names)
+
+        # Then we have general tags: pick any where prediction confidence > threshold
+        general_names = [labels[i] for i in self.general_indexes]
+
+        if general_mcut_enabled:
+            general_probs = np.array([x[1] for x in general_names])
+            general_thresh = mcut_threshold(general_probs)
+
+        general_res = [x for x in general_names if x[1] > general_thresh]
+        general_res = dict(general_res)
+
+        # Everything else is characters: pick any where prediction confidence > threshold
+        character_names = [labels[i] for i in self.character_indexes]
+
+        if character_mcut_enabled:
+            character_probs = np.array([x[1] for x in character_names])
+            character_thresh = mcut_threshold(character_probs)
+            character_thresh = max(0.15, character_thresh)
+
+        character_res = [x for x in character_names if x[1] > character_thresh]
+        character_res = dict(character_res)
+
+        sorted_general_strings = sorted(
+            general_res.items(),
+            key=lambda x: x[1],
+            reverse=True,
+        )
+        sorted_general_strings = [x[0] for x in sorted_general_strings]
+        sorted_general_strings = (
+            ", ".join(sorted_general_strings).replace("(", "\(").replace(")", "\)")
+        )
+
+        return sorted_general_strings, rating, character_res, general_res
+
+
+def tagger_ui():
+
+    predictor = Predictor()
+
+    dropdown_list = [
+        SWINV2_MODEL_DSV3_REPO,
+        CONV_MODEL_DSV3_REPO,
+        VIT_MODEL_DSV3_REPO,
+        VIT_LARGE_MODEL_DSV3_REPO,
+        EVA02_LARGE_MODEL_DSV3_REPO,
+        MOAT_MODEL_DSV2_REPO,
+        SWIN_MODEL_DSV2_REPO,
+        CONV_MODEL_DSV2_REPO,
+        CONV2_MODEL_DSV2_REPO,
+        VIT_MODEL_DSV2_REPO,
+    ]
+
+    with gr.Row():
+        with gr.Column(variant="panel"):
+            image = gr.Image(type="pil", image_mode="RGBA", label="Input")
+            model_repo = gr.Dropdown(
+                dropdown_list,
+                value=SWINV2_MODEL_DSV3_REPO,
+                label="模型",
+            )
+            with gr.Row():
+                general_thresh = gr.Slider(
+                    0,
+                    1,
+                    step=0.05,
+                    value=0.35,
+                    label="一般Tag阈值",
+                    scale=3,
+                )
+                general_mcut_enabled = gr.Checkbox(
+                    value=False,
+                    label="使用MCut阈值",
+                    scale=1,
+                )
+            with gr.Row():
+                character_thresh = gr.Slider(
+                    0,
+                    1,
+                    step=0.05,
+                    value=0.85,
+                    label="角色Tags阈值",
+                    scale=3,
+                )
+                character_mcut_enabled = gr.Checkbox(
+                    value=False,
+                    label="使用MCut阈值",
+                    scale=1,
+                )
+            with gr.Row():
+                clear = gr.ClearButton(
+                    components=[
+                        image,
+                        model_repo,
+                        general_thresh,
+                        general_mcut_enabled,
+                        character_thresh,
+                        character_mcut_enabled,
+                    ],
+                    variant="secondary",
+                    size="lg",
+                )
+                submit = gr.Button(value="提交", variant="primary", size="lg")
+        with gr.Column(variant="panel"):
+            sorted_general_strings = gr.Textbox(label="输出 (字符串)")
+            send_btn = gr.Button("发送到文生图",visible=False)
+            rating = gr.Label(label="分级")
+            character_res = gr.Label(label="输出 (角色)")
+            general_res = gr.Label(label="输出 (Tag)")
+            clear.add(
+                [
+                    sorted_general_strings,
+                    rating,
+                    character_res,
+                    general_res,
+                ]
+            )
+        sorted_general_strings.change(lambda s: gr.Button(visible=s is not None and len(s)), inputs=sorted_general_strings, outputs=send_btn)
+        submit.click(
+            predictor.predict,
+            inputs=[
+                image,
+                model_repo,
+                general_thresh,
+                general_mcut_enabled,
+                character_thresh,
+                character_mcut_enabled,
+            ],
+            outputs=[sorted_general_strings, rating, character_res, general_res],
+        )
+        return sorted_general_strings, send_btn