two more characters

README.md

@@ -1,12 +1,13 @@
 ---
-title: WaifuDiffusion Tagger
+title: WaifuDiffusion v1.4 Tags
 emoji: 💬
 colorFrom: blue
 colorTo: red
 sdk: gradio
-sdk_version: 4.39.0
+sdk_version: 3.16.2
 app_file: app.py
 pinned: false
+duplicated_from: NoCrypt/DeepDanbooru_string
 ---
 
 # Configuration
@@ -35,4 +36,4 @@ Path to your main application file (which contains either `gradio` or `streamlit
 Path is relative to the root of the repository.
 
 `pinned`: _boolean_
-Whether the Space stays on top of your list.
+Whether the Space stays on top of your list.

Utils/dbimutils.py

@@ -0,0 +1,54 @@
+# DanBooru IMage Utility functions
+
+import cv2
+import numpy as np
+from PIL import Image
+
+
+def smart_imread(img, flag=cv2.IMREAD_UNCHANGED):
+    if img.endswith(".gif"):
+        img = Image.open(img)
+        img = img.convert("RGB")
+        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+    else:
+        img = cv2.imread(img, flag)
+    return img
+
+
+def smart_24bit(img):
+    if img.dtype is np.dtype(np.uint16):
+        img = (img / 257).astype(np.uint8)
+
+    if len(img.shape) == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    elif img.shape[2] == 4:
+        trans_mask = img[:, :, 3] == 0
+        img[trans_mask] = [255, 255, 255, 255]
+        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
+    return img
+
+
+def make_square(img, target_size):
+    old_size = img.shape[:2]
+    desired_size = max(old_size)
+    desired_size = max(desired_size, target_size)
+
+    delta_w = desired_size - old_size[1]
+    delta_h = desired_size - old_size[0]
+    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
+    left, right = delta_w // 2, delta_w - (delta_w // 2)
+
+    color = [255, 255, 255]
+    new_im = cv2.copyMakeBorder(
+        img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
+    )
+    return new_im
+
+
+def smart_resize(img, size):
+    # Assumes the image has already gone through make_square
+    if img.shape[0] > size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
+    elif img.shape[0] < size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
+    return img

app.py

@@ -1,4 +1,8 @@
+from __future__ import annotations
+
 import argparse
+import functools
+import html
 import os
 
 import gradio as gr
@@ -6,56 +10,40 @@ import huggingface_hub
 import numpy as np
 import onnxruntime as rt
 import pandas as pd
-from PIL import Image
+import piexif
+import piexif.helper
+import PIL.Image
+
+from Utils import dbimutils
 
-TITLE = "WaifuDiffusion Tagger"
+TITLE = "WaifuDiffusion v1.4 Tags"
 DESCRIPTION = """
-Demo for the WaifuDiffusion tagger models
+Demo for:
+- [SmilingWolf/wd-v1-4-moat-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-moat-tagger-v2)
+- [SmilingWolf/wd-v1-4-swinv2-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-convnext-tagger-v2)
+- [SmilingWolf/wd-v1-4-convnext-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-convnext-tagger-v2)
+- [SmilingWolf/wd-v1-4-convnextv2-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-convnextv2-tagger-v2)
+- [SmilingWolf/wd-v1-4-vit-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger-v2)
+
+Includes "ready to copy" prompt and a prompt analyzer.
+
+Modified from [NoCrypt/DeepDanbooru_string](https://huggingface.co/spaces/NoCrypt/DeepDanbooru_string)  
+Modified from [hysts/DeepDanbooru](https://huggingface.co/spaces/hysts/DeepDanbooru)
+
+PNG Info code forked from [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
 
 Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
 """
 
-# 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
+HF_TOKEN = os.environ["HF_TOKEN"]
+MOAT_MODEL_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
+SWIN_MODEL_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
+CONV_MODEL_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
+CONV2_MODEL_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
+VIT_MODEL_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
 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 parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
@@ -66,273 +54,231 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-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
+def load_model(model_repo: str, model_filename: str) -> rt.InferenceSession:
+    path = huggingface_hub.hf_hub_download(
+        model_repo, model_filename, use_auth_token=HF_TOKEN
     )
-    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
+    model = rt.InferenceSession(path)
+    return model
 
-        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,
-            )
+def change_model(model_name):
+    global loaded_models
 
-        # Convert to numpy array
-        image_array = np.asarray(padded_image, dtype=np.float32)
+    if model_name == "MOAT":
+        model = load_model(MOAT_MODEL_REPO, MODEL_FILENAME)
+    elif model_name == "SwinV2":
+        model = load_model(SWIN_MODEL_REPO, MODEL_FILENAME)
+    elif model_name == "ConvNext":
+        model = load_model(CONV_MODEL_REPO, MODEL_FILENAME)
+    elif model_name == "ConvNextV2":
+        model = load_model(CONV2_MODEL_REPO, MODEL_FILENAME)
+    elif model_name == "ViT":
+        model = load_model(VIT_MODEL_REPO, MODEL_FILENAME)
 
-        # Convert PIL-native RGB to BGR
-        image_array = image_array[:, :, ::-1]
+    loaded_models[model_name] = model
+    return loaded_models[model_name]
 
-        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]
+def load_labels() -> list[str]:
+    path = huggingface_hub.hf_hub_download(
+        MOAT_MODEL_REPO, LABEL_FILENAME, use_auth_token=HF_TOKEN
+    )
+    df = pd.read_csv(path)
 
-        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)
+    tag_names = df["name"].tolist()
+    rating_indexes = list(np.where(df["category"] == 9)[0])
+    general_indexes = list(np.where(df["category"] == 0)[0])
+    character_indexes = list(np.where(df["category"] == 4)[0])
+    return tag_names, rating_indexes, general_indexes, character_indexes
 
-        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(")", "\)")
+def plaintext_to_html(text):
+    text = (
+        "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split("\n")]) + "</p>"
+    )
+    return text
+
+
+def predict(
+    image: PIL.Image.Image,
+    model_name: str,
+    general_threshold: float,
+    character_threshold: float,
+    tag_names: list[str],
+    rating_indexes: list[np.int64],
+    general_indexes: list[np.int64],
+    character_indexes: list[np.int64],
+):
+    global loaded_models
+
+    rawimage = image
+
+    model = loaded_models[model_name]
+    if model is None:
+        model = change_model(model_name)
+
+    _, height, width, _ = model.get_inputs()[0].shape
+
+    # Alpha to white
+    image = image.convert("RGBA")
+    new_image = PIL.Image.new("RGBA", image.size, "WHITE")
+    new_image.paste(image, mask=image)
+    image = new_image.convert("RGB")
+    image = np.asarray(image)
+
+    # PIL RGB to OpenCV BGR
+    image = image[:, :, ::-1]
+
+    image = dbimutils.make_square(image, height)
+    image = dbimutils.smart_resize(image, height)
+    image = image.astype(np.float32)
+    image = np.expand_dims(image, 0)
+
+    input_name = model.get_inputs()[0].name
+    label_name = model.get_outputs()[0].name
+    probs = model.run([label_name], {input_name: image})[0]
+
+    labels = list(zip(tag_names, probs[0].astype(float)))
+
+    # First 4 labels are actually ratings: pick one with argmax
+    ratings_names = [labels[i] for i in rating_indexes]
+    rating = dict(ratings_names)
+
+    # Then we have general tags: pick any where prediction confidence > threshold
+    general_names = [labels[i] for i in general_indexes]
+    general_res = [x for x in general_names if x[1] > general_threshold]
+    general_res = dict(general_res)
+
+    # Everything else is characters: pick any where prediction confidence > threshold
+    character_names = [labels[i] for i in character_indexes]
+    character_res = [x for x in character_names if x[1] > character_threshold]
+    character_res = dict(character_res)
+
+    b = dict(sorted(general_res.items(), key=lambda item: item[1], reverse=True))
+    a = (
+        ", ".join(list(b.keys()))
+        .replace("_", " ")
+        .replace("(", "\(")
+        .replace(")", "\)")
+    )
+    c = ", ".join(list(b.keys()))
+
+    items = rawimage.info
+    geninfo = ""
+
+    if "exif" in rawimage.info:
+        exif = piexif.load(rawimage.info["exif"])
+        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
+        try:
+            exif_comment = piexif.helper.UserComment.load(exif_comment)
+        except ValueError:
+            exif_comment = exif_comment.decode("utf8", errors="ignore")
+
+        items["exif comment"] = exif_comment
+        geninfo = exif_comment
+
+        for field in [
+            "jfif",
+            "jfif_version",
+            "jfif_unit",
+            "jfif_density",
+            "dpi",
+            "exif",
+            "loop",
+            "background",
+            "timestamp",
+            "duration",
+        ]:
+            items.pop(field, None)
+
+    geninfo = items.get("parameters", geninfo)
+
+    info = f"""
+<p><h4>PNG Info</h4></p>    
+"""
+    for key, text in items.items():
+        info += (
+            f"""
+<div>
+<p><b>{plaintext_to_html(str(key))}</b></p>
+<p>{plaintext_to_html(str(text))}</p>
+</div>
+""".strip()
+            + "\n"
         )
 
-        return sorted_general_strings, rating, character_res, general_res
+    if len(info) == 0:
+        message = "Nothing found in the image."
+        info = f"<div><p>{message}<p></div>"
+
+    return (a, c, rating, character_res, general_res, info)
 
 
 def main():
+    global loaded_models
+    loaded_models = {
+        "MOAT": None,
+        "SwinV2": None,
+        "ConvNext": None,
+        "ConvNextV2": None,
+        "ViT": None,
+    }
+
     args = parse_args()
 
-    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.Blocks(title=TITLE) as demo:
-        with gr.Column():
-            gr.Markdown(
-                value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>"
-            )
-            gr.Markdown(value=DESCRIPTION)
-            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="Model",
-                    )
-                    with gr.Row():
-                        general_thresh = gr.Slider(
-                            0,
-                            1,
-                            step=args.score_slider_step,
-                            value=args.score_general_threshold,
-                            label="General Tags Threshold",
-                            scale=3,
-                        )
-                        general_mcut_enabled = gr.Checkbox(
-                            value=False,
-                            label="Use MCut threshold",
-                            scale=1,
-                        )
-                    with gr.Row():
-                        character_thresh = gr.Slider(
-                            0,
-                            1,
-                            step=args.score_slider_step,
-                            value=args.score_character_threshold,
-                            label="Character Tags Threshold",
-                            scale=3,
-                        )
-                        character_mcut_enabled = gr.Checkbox(
-                            value=False,
-                            label="Use MCut threshold",
-                            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="Submit", variant="primary", size="lg")
-                with gr.Column(variant="panel"):
-                    sorted_general_strings = gr.Textbox(label="Output (string)")
-                    rating = gr.Label(label="Rating")
-                    character_res = gr.Label(label="Output (characters)")
-                    general_res = gr.Label(label="Output (tags)")
-                    clear.add(
-                        [
-                            sorted_general_strings,
-                            rating,
-                            character_res,
-                            general_res,
-                        ]
-                    )
-
-        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],
-        )
+    change_model("MOAT")
 
-        gr.Examples(
-            [["power.jpg", SWINV2_MODEL_DSV3_REPO, 0.35, False, 0.85, False]],
-            inputs=[
-                image,
-                model_repo,
-                general_thresh,
-                general_mcut_enabled,
-                character_thresh,
-                character_mcut_enabled,
-            ],
-        )
+    tag_names, rating_indexes, general_indexes, character_indexes = load_labels()
 
-    demo.queue(max_size=10)
-    demo.launch()
+    func = functools.partial(
+        predict,
+        tag_names=tag_names,
+        rating_indexes=rating_indexes,
+        general_indexes=general_indexes,
+        character_indexes=character_indexes,
+    )
+
+    gr.Interface(
+        fn=func,
+        inputs=[
+            gr.Image(type="pil", label="Input"),
+            gr.Radio(
+                ["MOAT", "SwinV2", "ConvNext", "ConvNextV2", "ViT"],
+                value="MOAT",
+                label="Model",
+            ),
+            gr.Slider(
+                0,
+                1,
+                step=args.score_slider_step,
+                value=args.score_general_threshold,
+                label="General Tags Threshold",
+            ),
+            gr.Slider(
+                0,
+                1,
+                step=args.score_slider_step,
+                value=args.score_character_threshold,
+                label="Character Tags Threshold",
+            ),
+        ],
+        outputs=[
+            gr.Textbox(label="Output (string)"),
+            gr.Textbox(label="Output (raw string)"),
+            gr.Label(label="Rating"),
+            gr.Label(label="Output (characters)"),
+            gr.Label(label="Output (tags)"),
+            gr.HTML(),
+        ],
+        examples=[["power.jpg", "MOAT", 0.35, 0.85]],
+        title=TITLE,
+        description=DESCRIPTION,
+        allow_flagging="never",
+    ).launch(
+        enable_queue=True,
+        share=args.share,
+    )
 
 
 if __name__ == "__main__":

requirements.txt

@@ -1,3 +1,5 @@
 pillow>=9.0.0
+piexif>=1.1.3
 onnxruntime>=1.12.0
+opencv-python
 huggingface-hub