p1atdev/wd-swinv2-tagger-v3-hf

Instalación

pip install transformers

Pipeline

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},
# ...

AutoModel

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]) # tomar los primeros logits

# obtener probabilidades
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 # umbral del 35%
}
print(results) # también se incluyen etiquetas de calificación y etiquetas de personajes
# {'1girl': tensor(0.9974),
#  'solo': tensor(0.9720),
#  'dress': tensor(0.9539),
#  'hat': tensor(0.9512),
#  'outdoors': tensor(0.9439),
# ...

Aceleración con Optimum

El uso de Optimum acelera el modelo aproximadamente un 30% y reduce el tamaño del modelo en un 50% en comparación con la versión de transformers, aunque la precisión se degrada ligeramente.

pip install optimum[onnxruntime]

-from transformers import pipeline
+from optimum.pipelines 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.9966088533401489},
#  {'label': 'solo', 'score': 0.9740601778030396},
#  {'label': 'dress', 'score': 0.9618403911590576},
#  {'label': 'hat', 'score': 0.9563733339309692},
#  {'label': 'outdoors', 'score': 0.945336639881134},
# ...

Etiquetas

• Todas las etiquetas de calificación tienen el prefijo rating:.
• Todas las etiquetas de personajes tienen el prefijo character:. Etiquetas de calificación: rating:general, rating:sensitive... Etiquetas de personajes: character:frieren, character:hatsune miku...