--- library_name: transformers license: apache-2.0 base_model: SmilingWolf/wd-swinv2-tagger-v3 inference: false tags: - wd-tagger --- # WD SwinV2 Tagger v3 with 🤗 transformers Converted from [SmilingWolf/wd-swinv2-tagger-v3](https://huggingface.co/SmilingWolf/wd-swinv2-tagger-v3) to transformers library format. ## Example ### Pipeline ```py from transformers import pipeline pipe = pipeline( "image-classification", model="p1atdev/wd-swinv2-tagger-v3-hf", trust_remote_code=True, ) print(pipe("sample.webp", top_k=15)) #[{'label': '1girl', 'score': 0.9973934888839722}, # {'label': 'solo', 'score': 0.9719744324684143}, # {'label': 'dress', 'score': 0.9539461135864258}, # {'label': 'hat', 'score': 0.9511678218841553}, # {'label': 'outdoors', 'score': 0.9438753128051758}, # {'label': 'sky', 'score': 0.9195725917816162}, # {'label': 'sitting', 'score': 0.9178725481033325}, # {'label': 'looking up', 'score': 0.9122412800788879}, # {'label': 'short hair', 'score': 0.8630313873291016}, # {'label': 'cloud', 'score': 0.8609118461608887}, # {'label': 'brown hair', 'score': 0.7723952531814575}, # {'label': 'short sleeves', 'score': 0.7649227380752563}, # {'label': 'day', 'score': 0.7641971111297607}, # {'label': 'rating:general', 'score': 0.7605368494987488}, # {'label': 'white dress', 'score': 0.7596388459205627}] ``` ### AutoModel ```py from PIL import Image import numpy as np import torch from transformers import ( AutoImageProcessor, AutoModelForImageClassification, ) MODEL_NAME = "p1atdev/wd-swinv2-tagger-v3-hf" model = AutoModelForImageClassification.from_pretrained( MODEL_NAME, ) processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True) image = Image.open("sample.webp") inputs = processor.preprocess(image, return_tensors="pt") with torch.no_grad(): outputs = model(**inputs.to(model.device, model.dtype)) logits = torch.sigmoid(outputs.logits[0]) # take the first logits # get probabilities results = {model.config.id2label[i]: logit.float() for i, logit in enumerate(logits)} results = { k: v for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True) if v > 0.35 # 35% threshold } print(results) # rating tags and character tags are also included #{'1girl': tensor(0.9974), # 'solo': tensor(0.9720), # 'dress': tensor(0.9539), # 'hat': tensor(0.9512), # 'outdoors': tensor(0.9439), # 'sky': tensor(0.9196), # 'sitting': tensor(0.9179), # 'looking up': tensor(0.9122), # 'short hair': tensor(0.8630), # 'cloud': tensor(0.8609), # 'brown hair': tensor(0.7724), # 'short sleeves': tensor(0.7649), # 'day': tensor(0.7642), # 'rating:general': tensor(0.7605), # ... ``` ## Labels All of rating tags have prefix `rating:` and character tags have prefix `character:`. - Rating tags: `rating:general`, `rating:sensitive`, ... - Character tags: `character:frieren`, `character:hatsune miku`, ...