File size: 2,240 Bytes
74bac1e
 
01e8564
 
 
624818f
 
74bac1e
 
01e8564
 
 
 
 
 
e4465bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
01e8564
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
624818f
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
---
library_name: transformers
license: apache-2.0
base_model: SmilingWolf/wd-swinv2-tagger-v3
inference: false
tags:
- wd-v14-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=20))
#[{'label': '1girl', 'score': 0.996832549571991},
# {'label': 'solo', 'score': 0.958362340927124},
# {'label': 'dress', 'score': 0.9418023228645325},
# {'label': 'hat', 'score': 0.9263725280761719},
# {'label': 'sitting', 'score': 0.9178062677383423},
# {'label': 'looking_up', 'score': 0.8978123068809509},
# {'label': 'short_hair', 'score': 0.8242536187171936},
# {'label': 'sky', 'score': 0.784591794013977},
# {'label': 'outdoors', 'score': 0.767611563205719},
# {'label': 'rating:general', 'score': 0.7561964988708496}]
```

### 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,
    torch_dtype=torch.bfloat16,
)
model.eval()
processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)

image = Image.open("sample.webp")
inputs = processor.preprocess(image, return_tensors="pt")

outputs = model(**inputs.to(model.device, model.dtype))
logits = torch.sigmoid(outputs.logits[0])

# 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)
}
print(results)  # rating tags and character tags are also included
#{'1girl': tensor(0.9968),
# 'solo': tensor(0.9584),
# 'dress': tensor(0.9418),
# 'hat': tensor(0.9264),
# 'sitting': tensor(0.9178),
# 'looking_up': tensor(0.8978),
# 'short_hair': tensor(0.8243),
# 'sky': tensor(0.7846),
# 'outdoors': tensor(0.7676),
# 'rating:general': tensor(0.7562),
# ...
```