server/tagger.py

import numpy as np
import onnxruntime as rt
import pandas as pd
from PIL import Image
import io
import base64
import huggingface_hub

MODEL_FILENAME = "model.onnx"
LABEL_FILENAME = "selected_tags.csv"

KAOMOJIS = {
    "0_0", "(o)_(o)", "+_+", "+_-", "._.", "<o>_<o>", "<|>_<|>",
    "=_=", ">_<", "3_3", "6_9", ">_o", "@_@", "^_^", "o_o",
    "u_u", "x_x", "|_|", "||_||",
}


def load_labels(dataframe):
    name_series = dataframe["name"].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):
    sorted_probs = probs[probs.argsort()[::-1]]
    difs = sorted_probs[:-1] - sorted_probs[1:]
    t = difs.argmax()
    return (sorted_probs[t] + sorted_probs[t + 1]) / 2


class TaggerPredictor:
    def __init__(self):
        self.model_target_size = None
        self.last_loaded_repo = None
        self.model = None
        self.tag_names = []
        self.rating_indexes = []
        self.general_indexes = []
        self.character_indexes = []

    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]

        self.model = rt.InferenceSession(model_path)
        _, height, width, _ = self.model.get_inputs()[0].shape
        self.model_target_size = height
        self.last_loaded_repo = model_repo

    def prepare_image(self, image: Image.Image):
        target_size = self.model_target_size

        canvas = Image.new("RGBA", image.size, (255, 255, 255))
        canvas.alpha_composite(image.convert("RGBA"))
        image = canvas.convert("RGB")

        max_dim = max(image.size)
        pad_left = (max_dim - image.size[0]) // 2
        pad_top = (max_dim - image.size[1]) // 2

        padded = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
        padded.paste(image, (pad_left, pad_top))

        if max_dim != target_size:
            padded = padded.resize((target_size, target_size), Image.BICUBIC)

        arr = np.asarray(padded, dtype=np.float32)
        arr = arr[:, :, ::-1]  # RGB -> BGR
        return np.expand_dims(arr, axis=0)

    def predict(self, image_b64, model_repo, general_thresh, general_mcut,
                character_thresh, character_mcut):
        self.load_model(model_repo)

        img_bytes = base64.b64decode(image_b64)
        image = Image.open(io.BytesIO(img_bytes))
        prepared = 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: prepared})[0]

        labels = list(zip(self.tag_names, preds[0].astype(float)))

        # Ratings
        ratings = {labels[i][0]: float(labels[i][1]) for i in self.rating_indexes}

        # General tags
        general_names = [labels[i] for i in self.general_indexes]
        if general_mcut:
            probs = np.array([x[1] for x in general_names])
            general_thresh = float(mcut_threshold(probs))
        general_res = {n: float(s) for n, s in general_names if s > general_thresh}

        # Characters
        character_names = [labels[i] for i in self.character_indexes]
        if character_mcut:
            probs = np.array([x[1] for x in character_names])
            character_thresh = max(0.15, float(mcut_threshold(probs)))
        character_res = {n: float(s) for n, s in character_names if s > character_thresh}

        # Sorted string
        sorted_tags = sorted(general_res.items(), key=lambda x: x[1], reverse=True)
        tags_string = ", ".join(
            x[0].replace("(", "\\(").replace(")", "\\)")
            for x in sorted_tags
        )

        return {
            "tags": tags_string,
            "rating": ratings,
            "characters": character_res,
            "general": general_res,
        }


# Singleton
tagger_predictor = TaggerPredictor()