File size: 2,908 Bytes
74bac1e 01e8564 624818f db283a4 74bac1e 01e8564 e4465bd bfb6030 e4465bd 8563064 e4465bd 01e8564 bfb6030 01e8564 ed14c24 22b2730 01e8564 22b2730 01e8564 5e60fa0 01e8564 d6fd8d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
---
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`, ...
|