imgutils.tagging.wd14
- Overview:
This module provides utilities for image tagging using WD14 taggers. It includes functions for loading models, processing images, and extracting tags.
The module is inspired by the SmilingWolf/wd-v1-4-tags project on Hugging Face.
convert_wd14_emb_to_prediction
- imgutils.tagging.wd14.convert_wd14_emb_to_prediction(emb: ndarray, model_name: str = 'SwinV2_v3', general_threshold: float = 0.35, general_mcut_enabled: bool = False, character_threshold: float = 0.85, character_mcut_enabled: bool = False, no_underline: bool = False, drop_overlap: bool = False, fmt=('rating', 'general', 'character'))[source]
Convert WD14 embedding to understandable prediction result.
- Parameters:
emb (numpy.ndarray) – The 1-dim extracted embedding.
model_name (str) – The name of the model to use.
general_threshold (float) – The threshold for general tags.
general_mcut_enabled (bool) – If True, applies MCut thresholding to general tags.
character_threshold (float) – The threshold for character tags.
character_mcut_enabled (bool) – If True, applies MCut thresholding to character tags.
no_underline (bool) – If True, replaces underscores in tag names with spaces.
drop_overlap (bool) – If True, drops overlapping tags.
fmt – Return format, default is
('rating', 'general', 'character')
.
- Returns:
Prediction result based on the provided fmt.
Note
Only the embeddings not get normalized can be converted to understandable prediction result.
- Example:
>>> import os >>> from imgutils.tagging import get_wd14_tags, convert_wd14_emb_to_prediction >>> >>> # extract the feature embedding >>> embedding = get_wd14_tags('skadi.jpg', fmt='embedding') >>> >>> # convert to understandable result >>> rating, general, character = convert_wd14_emb_to_prediction(embedding) >>> # these 3 dicts will be the same as that returned by `get_wd14_tags('skadi.jpg')`