Anime Tagger swinv2_base_window8_256.e621v1-full

Model Details

• Model Type: Multilabel Image classification / feature backbone
• Model Stats:
  • Params: 95.9M
  • FLOPs / MACs: 121.6G / 60.7G
  • Image size: train = 448 x 448, test = 448 x 448
• Dataset: animetimm/e621-wdtagger-v1-w640-ws-full
  • Tags Count: 8783
    • General (#0) Tags Count: 7180
    • Character (#4) Tags Count: 494
    • Species (#5) Tags Count: 1106
    • Rating (#9) Tags Count: 3

Results

• Macro/Micro@0.40 means the metrics on the threshold 0.40.
• Macro@Best means the mean metrics on the tag-level thresholds on each tags, which should have the best F1 scores.

Thresholds

• Micro@Thr means the metrics on the category-level suggested thresholds, which are listed in the table above.
• Macro@0.40 means the metrics on the threshold 0.40.
• Macro@Best means the metrics on the tag-level thresholds on each tags, which should have the best F1 scores.

For tag-level thresholds, you can find them in selected_tags.csv.

How to Use

We provided a sample image for our code samples, you can find it here.

Use TIMM And Torch

Install dghs-imgutils, timm and other necessary requirements with the following command

pip install 'dghs-imgutils>=0.17.0' torch huggingface_hub timm pillow pandas

After that you can load this model with timm library, and use it for train, validation and test, with the following code

import json

import pandas as pd
import torch
from huggingface_hub import hf_hub_download
from imgutils.data import load_image
from imgutils.preprocess import create_torchvision_transforms
from timm import create_model

repo_id = 'animetimm/swinv2_base_window8_256.e621v1-full'
model = create_model(f'hf-hub:{repo_id}', pretrained=True)
model.eval()

with open(hf_hub_download(repo_id=repo_id, repo_type='model', filename='preprocess.json'), 'r') as f:
    preprocessor = create_torchvision_transforms(json.load(f)['test'])
# Compose(
#     PadToSize(size=(512, 512), interpolation=bilinear, background_color=white)
#     Resize(size=448, interpolation=bicubic, max_size=None, antialias=True)
#     CenterCrop(size=[448, 448])
#     MaybeToTensor()
#     Normalize(mean=tensor([0.4850, 0.4560, 0.4060]), std=tensor([0.2290, 0.2240, 0.2250]))
# )

image = load_image('https://huggingface.co/animetimm/swinv2_base_window8_256.e621v1-full/resolve/main/sample.webp')
input_ = preprocessor(image).unsqueeze(0)
# input_, shape: torch.Size([1, 3, 448, 448]), dtype: torch.float32
with torch.no_grad():
    output = model(input_)
    prediction = torch.sigmoid(output)[0]
# output, shape: torch.Size([1, 8783]), dtype: torch.float32
# prediction, shape: torch.Size([8783]), dtype: torch.float32

df_tags = pd.read_csv(
    hf_hub_download(repo_id=repo_id, repo_type='model', filename='selected_tags.csv'),
    keep_default_na=False
)
tags = df_tags['name']
mask = prediction.numpy() >= df_tags['best_threshold']
print(dict(zip(tags[mask].tolist(), prediction[mask].tolist())))
# {'explicit': 0.9982947707176208,
#  'anthro': 0.9554250240325928,
#  'female': 0.9985587000846863,
#  'solo': 0.9950961470603943,
#  'genitals': 0.9889377355575562,
#  'clothing': 0.883478581905365,
#  'breasts': 0.9927160143852234,
#  'nipples': 0.9559563994407654,
#  'simple_background': 0.5567265748977661,
#  'pussy': 0.9784653782844543,
#  'big_breasts': 0.9587277770042419,
#  'horn': 0.8108257055282593,
#  'thick_thighs': 0.7229651212692261,
#  'navel': 0.6964868903160095,
#  'areola': 0.5498285889625549,
#  'huge_breasts': 0.7597795724868774,
#  'wide_hips': 0.6507709622383118,
#  'legwear': 0.8271871209144592,
#  'red_eyes': 0.3962002992630005,
#  'handwear': 0.36074572801589966,
#  'gloves': 0.4130750298500061,
#  'slightly_chubby': 0.6015251278877258,
#  'scales': 0.3458760380744934,
#  'non-mammal_breasts': 0.8451875448226929,
#  'curvy_figure': 0.5748545527458191,
#  'mostly_nude': 0.41494160890579224,
#  'stockings': 0.4295252859592438,
#  'clitoris': 0.3801887035369873,
#  'armwear': 0.49424344301223755,
#  'voluptuous': 0.6553569436073303,
#  'non-mammal_nipples': 0.7104668021202087,
#  'big_areola': 0.5583626627922058,
#  'big_nipples': 0.6151208281517029,
#  'navel_piercing': 0.4667469561100006,
#  'huge_nipples': 0.26599329710006714,
#  'ring_(jewelry)': 0.2037716656923294,
#  'scalie': 0.9633249044418335,
#  'argonian': 0.853269636631012}

Use ONNX Model For Inference

Install dghs-imgutils with the following command

pip install 'dghs-imgutils>=0.17.0'

Use multilabel_timm_predict function with the following code

from imgutils.generic import multilabel_timm_predict

general, character, species, rating = multilabel_timm_predict(
    'https://huggingface.co/animetimm/swinv2_base_window8_256.e621v1-full/resolve/main/sample.webp',
    repo_id='animetimm/swinv2_base_window8_256.e621v1-full',
    fmt=('general', 'character', 'species', 'rating'),
)

print(general)
# {'female': 0.9985586404800415,
#  'solo': 0.9950962066650391,
#  'breasts': 0.9927159547805786,
#  'genitals': 0.9889377355575562,
#  'pussy': 0.9784653782844543,
#  'big_breasts': 0.9587278366088867,
#  'nipples': 0.9559563398361206,
#  'anthro': 0.9554250836372375,
#  'clothing': 0.8834788203239441,
#  'non-mammal_breasts': 0.8451874256134033,
#  'legwear': 0.8271874189376831,
#  'horn': 0.8108260035514832,
#  'huge_breasts': 0.7597798705101013,
#  'thick_thighs': 0.7229658961296082,
#  'non-mammal_nipples': 0.7104665637016296,
#  'navel': 0.696487545967102,
#  'voluptuous': 0.6553572416305542,
#  'wide_hips': 0.65077143907547,
#  'big_nipples': 0.6151217818260193,
#  'slightly_chubby': 0.6015253067016602,
#  'curvy_figure': 0.5748555064201355,
#  'big_areola': 0.5583634972572327,
#  'simple_background': 0.5567269325256348,
#  'areola': 0.5498292446136475,
#  'armwear': 0.4942445158958435,
#  'navel_piercing': 0.466747522354126,
#  'stockings': 0.42952579259872437,
#  'mostly_nude': 0.4149421453475952,
#  'gloves': 0.4130757451057434,
#  'red_eyes': 0.3962007761001587,
#  'clitoris': 0.3801894187927246,
#  'handwear': 0.36074620485305786,
#  'scales': 0.34587666392326355,
#  'huge_nipples': 0.26599401235580444,
#  'ring_(jewelry)': 0.20377254486083984}
print(character)
# {}
print(species)
# {'scalie': 0.963325023651123, 'argonian': 0.8532700538635254}
print(rating)
# {'explicit': 0.9982947111129761}

For further information, see documentation of function multilabel_timm_predict.