From fafc9e12072ca73ee9e84414f2d420c7a076fd12 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Mon, 8 Sep 2025 17:09:52 +0800 Subject: [PATCH 01/32] dev(narugo): add pixai tagger x, ci skip --- zoo/pixai_tagger/__init__.py | 0 zoo/pixai_tagger/export.py | 87 +++++++++++++ zoo/pixai_tagger/min_script.py | 228 +++++++++++++++++++++++++++++++++ zoo/pixai_tagger/tags.py | 82 ++++++++++++ 4 files changed, 397 insertions(+) create mode 100644 zoo/pixai_tagger/__init__.py create mode 100644 zoo/pixai_tagger/export.py create mode 100644 zoo/pixai_tagger/min_script.py create mode 100644 zoo/pixai_tagger/tags.py diff --git a/zoo/pixai_tagger/__init__.py b/zoo/pixai_tagger/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py new file mode 100644 index 00000000000..4b84f5dd931 --- /dev/null +++ b/zoo/pixai_tagger/export.py @@ -0,0 +1,87 @@ +import json +import os.path + +import numpy as np +import pandas as pd +from ditk import logging +from hbutils.system import TemporaryDirectory +from hfutils.operate import get_hf_client + +from imgutils.preprocess import parse_torchvision_transforms +from zoo.pixai_tagger.tags import load_tags +from .min_script import EndpointHandler + + +def sync(src_repo: str, dst_repo: str): + hf_client = get_hf_client() + if not hf_client.repo_exists(repo_id=dst_repo, repo_type='model'): + hf_client.create_repo(repo_id=dst_repo, repo_type='model', private=True) + + handler = EndpointHandler(repo_id=src_repo) + with TemporaryDirectory() as upload_dir: + preprocessor = handler.transform + preprocessor_file = os.path.join(upload_dir, 'preprocessor.json') + logging.info(f'Dumping preprocessor:\n{preprocessor}\nto file {preprocessor_file!r}.') + with open(preprocessor_file, 'w') as f: + json.dump({ + 'stages': parse_torchvision_transforms(handler.transform), + }, f, sort_keys=True, ensure_ascii=False, indent=4) + + logging.info('Scanning tags ...') + categories = np.zeros((len(handler.index_to_tag_map),), dtype=np.int32) + idx = np.array(range(len(categories))) + categories[idx < handler.gen_tag_count] = 0 + categories[idx >= handler.gen_tag_count] = 4 + df_src_tags = pd.DataFrame({ + 'name': [v for _, v in sorted(handler.index_to_tag_map.items())], + 'category': categories, + }) + df_tags = load_tags(df_src_tags) + exts = [] + for titem in df_tags.to_dict('records'): + if titem['category'] == 4: + if titem['name'] in handler.character_ip_mapping: + exts.append(handler.character_ip_mapping[titem['name']]) + else: + exts.append([]) + else: + exts.append([]) + df_tags['ips'] = exts + logging.info(f'Tags:\n{df_tags}') + df_tags.to_csv(os.path.join(upload_dir, 'selected_tags.csv'), index=False) + + with open(os.path.join(upload_dir, 'categories.json'), 'w') as f: + json.dump([ + { + "category": 0, + "name": "general" + }, + { + "category": 4, + "name": "character" + }, + ], f, sort_keys=True, ensure_ascii=False, indent=4) + df_th = pd.DataFrame([ + {'category': 0, 'name': 'general', 'threshold': handler.default_general_threshold}, + {'category': 4, 'name': 'character', 'threshold': handler.default_character_threshold}, + ]) + df_th.to_csv(os.path.join(upload_dir, 'thresholds.csv'), index=False) + + handler.model + + os.system(f'tree {upload_dir!r}') + input() + # print(df_tags) + + # pprint(handler.index_to_tag_map) + # pprint(handler.gen_tag_count) + # pprint(handler.character_tag_count) + # pprint(handler.character_ip_mapping) + + +if __name__ == '__main__': + logging.try_init_root(logging.INFO) + sync( + src_repo='pixai-labs/pixai-tagger-v0.9', + dst_repo='deepghs/pixai-tagger-v0.9-onnx' + ) diff --git a/zoo/pixai_tagger/min_script.py b/zoo/pixai_tagger/min_script.py new file mode 100644 index 00000000000..40cf338178d --- /dev/null +++ b/zoo/pixai_tagger/min_script.py @@ -0,0 +1,228 @@ +import base64 +import io +import json +import logging +import time +from pathlib import Path +from typing import Any + +import requests +import timm +import torch +import torchvision.transforms as transforms +from PIL import Image +from huggingface_hub import hf_hub_download + +from imgutils.preprocess import parse_torchvision_transforms + + +class TaggingHead(torch.nn.Module): + def __init__(self, input_dim, num_classes): + super().__init__() + self.input_dim = input_dim + self.num_classes = num_classes + self.head = torch.nn.Sequential(torch.nn.Linear(input_dim, num_classes)) + + def forward(self, x): + logits = self.head(x) + probs = torch.nn.functional.sigmoid(logits) + return probs + + +def get_tags(tags_file: Path) -> tuple[dict[str, int], int, int]: + with tags_file.open("r", encoding="utf-8") as f: + tag_info = json.load(f) + tag_map = tag_info["tag_map"] + tag_split = tag_info["tag_split"] + gen_tag_count = tag_split["gen_tag_count"] + character_tag_count = tag_split["character_tag_count"] + return tag_map, gen_tag_count, character_tag_count + + +def get_character_ip_mapping(mapping_file: Path): + with mapping_file.open("r", encoding="utf-8") as f: + mapping = json.load(f) + return mapping + + +def get_encoder(): + base_model_repo = "hf_hub:SmilingWolf/wd-eva02-large-tagger-v3" + encoder = timm.create_model(base_model_repo, pretrained=False) + encoder.reset_classifier(0) + return encoder + + +def get_decoder(): + decoder = TaggingHead(1024, 13461) + return decoder + + +def get_model(): + encoder = get_encoder() + decoder = get_decoder() + model = torch.nn.Sequential(encoder, decoder) + return model + + +def load_model(weights_file, device): + model = get_model() + states_dict = torch.load(weights_file, map_location=device, weights_only=True) + model.load_state_dict(states_dict) + model.to(device) + model.eval() + return model + + +def pure_pil_alpha_to_color_v2( + image: Image.Image, color: tuple[int, int, int] = (255, 255, 255) +) -> Image.Image: + """ + Convert a PIL image with an alpha channel to a RGB image. + This is a workaround for the fact that the model expects a RGB image, but the image may have an alpha channel. + This function will convert the image to a RGB image, and fill the alpha channel with the given color. + The alpha channel is the 4th channel of the image. + """ + image.load() # needed for split() + background = Image.new("RGB", image.size, color) + background.paste(image, mask=image.split()[3]) # 3 is the alpha channel + return background + + +def pil_to_rgb(image: Image.Image) -> Image.Image: + if image.mode == "RGBA": + image = pure_pil_alpha_to_color_v2(image) + elif image.mode == "P": + image = pure_pil_alpha_to_color_v2(image.convert("RGBA")) + else: + image = image.convert("RGB") + return image + + +class EndpointHandler: + def __init__(self, repo_id: str = 'pixai-labs/pixai-tagger-v0.9'): + weights_file = Path(hf_hub_download( + repo_id=repo_id, + repo_type='model', + filename="model_v0.9.pth", + )) + tags_file = Path(hf_hub_download( + repo_id=repo_id, + repo_type='model', + filename="tags_v0.9_13k.json", + )) + mapping_file = Path(hf_hub_download( + repo_id=repo_id, + repo_type='model', + filename="char_ip_map.json", + )) + + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.model = load_model(str(weights_file), self.device) + self.transform = transforms.Compose( + [ + transforms.Resize((448, 448)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), + ] + ) + self.fetch_image_timeout = 5.0 + self.default_general_threshold = 0.3 + self.default_character_threshold = 0.85 + + tag_map, self.gen_tag_count, self.character_tag_count = get_tags(tags_file) + + # Invert the tag_map for efficient index-to-tag lookups + self.index_to_tag_map = {v: k for k, v in tag_map.items()} + + self.character_ip_mapping = get_character_ip_mapping(mapping_file) + + def __call__(self, data: dict[str, Any]) -> dict[str, Any]: + inputs = data.pop("inputs", data) + + fetch_start_time = time.time() + if isinstance(inputs, Image.Image): + image = inputs + elif image_url := inputs.pop("url", None): + with requests.get( + image_url, stream=True, timeout=self.fetch_image_timeout + ) as res: + res.raise_for_status() + image = Image.open(res.raw) + elif image_base64_encoded := inputs.pop("image", None): + image = Image.open(io.BytesIO(base64.b64decode(image_base64_encoded))) + else: + raise ValueError(f"No image or url provided: {data}") + # remove alpha channel if it exists + image = pil_to_rgb(image) + fetch_time = time.time() - fetch_start_time + + parameters = data.pop("parameters", {}) + general_threshold = parameters.pop( + "general_threshold", self.default_general_threshold + ) + character_threshold = parameters.pop( + "character_threshold", self.default_character_threshold + ) + + inference_start_time = time.time() + with torch.inference_mode(): + # Preprocess image on CPU, then pin memory for faster async transfer + image_tensor = self.transform(image).unsqueeze(0).pin_memory() + + # Asynchronously move image to GPU + image_tensor = image_tensor.to(self.device, non_blocking=True) + + # Run model on GPU + probs = self.model(image_tensor)[0] # Get probs for the single image + + # Perform thresholding directly on the GPU + general_mask = probs[: self.gen_tag_count] > general_threshold + character_mask = probs[self.gen_tag_count:] > character_threshold + + # Get the indices of positive tags on the GPU + general_indices = general_mask.nonzero(as_tuple=True)[0] + character_indices = ( + character_mask.nonzero(as_tuple=True)[0] + self.gen_tag_count + ) + + # Combine indices and move the small result tensor to the CPU + combined_indices = torch.cat((general_indices, character_indices)).cpu() + + inference_time = time.time() - inference_start_time + + post_process_start_time = time.time() + + cur_gen_tags = [] + cur_char_tags = [] + + # Use the efficient pre-computed map for lookups + for i in combined_indices: + idx = i.item() + tag = self.index_to_tag_map[idx] + if idx < self.gen_tag_count: + cur_gen_tags.append(tag) + else: + cur_char_tags.append(tag) + + ip_tags = [] + for tag in cur_char_tags: + if tag in self.character_ip_mapping: + ip_tags.extend(self.character_ip_mapping[tag]) + ip_tags = sorted(set(ip_tags)) + post_process_time = time.time() - post_process_start_time + + logging.info( + f"Timing - Fetch: {fetch_time:.3f}s, Inference: {inference_time:.3f}s, Post-process: {post_process_time:.3f}s, Total: {fetch_time + inference_time + post_process_time:.3f}s" + ) + + return { + "feature": cur_gen_tags, + "character": cur_char_tags, + "ip": ip_tags, + } + + +if __name__ == '__main__': + handler = EndpointHandler() + print(handler.transform) + print(parse_torchvision_transforms(handler.transform)) diff --git a/zoo/pixai_tagger/tags.py b/zoo/pixai_tagger/tags.py new file mode 100644 index 00000000000..3eabe7c1c7a --- /dev/null +++ b/zoo/pixai_tagger/tags.py @@ -0,0 +1,82 @@ +from functools import lru_cache + +import pandas as pd +from ditk import logging +from hfutils.operate import get_hf_client +from tqdm import tqdm +from waifuc.utils import srequest + +from zoo.wd14.tags import _get_tag_by_name, _db_session + +_CATEGORY_MAPS = { + 'general': 0, + 'character': 4, +} + + +@lru_cache() +def _get_rating_count_by_name(tag_name: str): + session = _db_session() + logging.info(f'Getting count for {tag_name!r} ...') + vs = srequest( + session, 'GET', f'https://danbooru.donmai.us/counts/posts.json', + params={'tags': f'rating:{tag_name}'} + ).json() + logging.info(f'Result of {tag_name!r}: {vs!r}') + return vs['counts']['posts'] + + +def load_tags(df_src_tags: pd.DataFrame): + hf_client = get_hf_client() + df_p_tags = pd.read_csv(hf_client.hf_hub_download( + repo_id='deepghs/site_tags', + repo_type='dataset', + filename='danbooru.donmai.us/tags.csv' + )) + logging.info(f'Loaded danbooru tags pool, columns: {df_p_tags.columns!r}') + d_p_tags = {(item['category'], item['name']): item for item in df_p_tags.to_dict('records')} + + rows = [] + src_tags = df_src_tags.to_dict('records') + for i in tqdm(range(len(src_tags)), desc='Scan Tags'): + tag_name = src_tags[i]['name'] + category = src_tags[i]['category'] + if (category, tag_name) in d_p_tags: + tag_id = d_p_tags[(category, tag_name)]['id'] + count = d_p_tags[(category, tag_name)]['post_count'] + elif category < 9: + logging.warning(f'Cannot find tag {tag_name!r}, category: {category!r}.') + tag_info = _get_tag_by_name(tag_name) + if tag_info['name'] != tag_name: + logging.warning(f'Not found matching tags for {tag_name!r}, will be ignored.') + tag_id = -1 + count = -1 + else: + logging.info(f'Tag info found from danbooru - {tag_info!r}.') + tag_id = tag_info['id'] + count = tag_info['post_count'] + if category != tag_info['category']: + logging.warning(f'Category not match for tag {tag_name!r}, ' + f'replace category {category!r} --> {tag_info["category"]!r}') + category = tag_info['category'] + else: + logging.warning(f'Unknown tag {tag_name!r} ...') + tag_id = -1 + count = -1 + + rows.append({ + 'id': i, + 'tag_id': tag_id, + 'name': tag_name, + 'category': category, + 'count': count, + }) + + df = pd.DataFrame(rows) + return df + + +if __name__ == '__main__': + logging.try_init_root(level=logging.INFO) + df = load_tags() + df.to_csv('test_df.csv', index=False) From 96ba7f920031a0639466c982383fa6e99415297a Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Mon, 8 Sep 2025 17:30:37 +0800 Subject: [PATCH 02/32] dev(narugo): add more --- zoo/pixai_tagger/export.py | 53 ++++++++++++++++++++++++++++++++++++-- zoo/pixai_tagger/onnx.py | 48 ++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 2 deletions(-) create mode 100644 zoo/pixai_tagger/onnx.py diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 4b84f5dd931..31c13a33d93 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -2,22 +2,29 @@ import os.path import numpy as np +import onnx +import onnxruntime import pandas as pd +import torch from ditk import logging from hbutils.system import TemporaryDirectory from hfutils.operate import get_hf_client +from imgutils.data import load_image from imgutils.preprocess import parse_torchvision_transforms from zoo.pixai_tagger.tags import load_tags +from zoo.utils import onnx_optimize, get_testfile from .min_script import EndpointHandler +from .onnx import get_model -def sync(src_repo: str, dst_repo: str): +def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): hf_client = get_hf_client() if not hf_client.repo_exists(repo_id=dst_repo, repo_type='model'): hf_client.create_repo(repo_id=dst_repo, repo_type='model', private=True) handler = EndpointHandler(repo_id=src_repo) + with TemporaryDirectory() as upload_dir: preprocessor = handler.transform preprocessor_file = os.path.join(upload_dir, 'preprocessor.json') @@ -67,7 +74,49 @@ def sync(src_repo: str, dst_repo: str): ]) df_th.to_csv(os.path.join(upload_dir, 'thresholds.csv'), index=False) - handler.model + dummy_image = load_image(get_testfile('6125785.jpg'), mode='RGB', force_background='white') + dummy_input = handler.transform(dummy_image).unsqueeze(0) + wrapped_model, (conv_features, _) = get_model(handler.model, dummy_input) + onnx_filename = os.path.join(upload_dir, 'model.onnx') + with TemporaryDirectory() as td: + temp_model_onnx = os.path.join(td, 'model.onnx') + logging.info(f'Exporting temporary ONNX model to {temp_model_onnx!r} ...') + torch.onnx.export( + wrapped_model, + dummy_input, + temp_model_onnx, + input_names=['input'], + output_names=['embedding', 'output'], + dynamic_axes={ + 'input': {0: 'batch_size'}, + 'embedding': {0: 'batch_size'}, + 'output': {0: 'batch_size'}, + }, + opset_version=14, + do_constant_folding=True, + export_params=True, + verbose=False, + custom_opsets=None, + ) + + model = onnx.load(temp_model_onnx) + if not no_optimize: + logging.info('Optimizing onnx model ...') + model = onnx_optimize(model) + + output_model_dir, _ = os.path.split(onnx_filename) + if output_model_dir: + os.makedirs(output_model_dir, exist_ok=True) + logging.info(f'Complete model saving to {onnx_filename!r} ...') + onnx.save(model, onnx_filename) + + session = onnxruntime.InferenceSession(onnx_filename) + o_embeddings, = session.run(['embedding'], {'input': dummy_input.numpy()}) + emb_1 = o_embeddings / np.linalg.norm(o_embeddings, axis=-1, keepdims=True) + emb_2 = conv_features.numpy() / np.linalg.norm(conv_features.numpy(), axis=-1, keepdims=True) + emb_sims = (emb_1 * emb_2).sum() + logging.info(f'Similarity of the embeddings is {emb_sims:.5f}.') + assert emb_sims >= 0.98, f'Similarity of the embeddings is {emb_sims:.5f}, ONNX validation failed.' os.system(f'tree {upload_dir!r}') input() diff --git a/zoo/pixai_tagger/onnx.py b/zoo/pixai_tagger/onnx.py new file mode 100644 index 00000000000..759b9b7aaf7 --- /dev/null +++ b/zoo/pixai_tagger/onnx.py @@ -0,0 +1,48 @@ +import logging + +import torch +from torch import nn + + +class ModuleWrapper(nn.Module): + def __init__(self, base_module: nn.Module, classifier: nn.Module): + super().__init__() + self.base_module = base_module + self.classifier = classifier + + self._output_features = None + self._register_hook() + + def _register_hook(self): + def hook_fn(module, input_tensor, output_tensor): + assert isinstance(input_tensor, tuple) and len(input_tensor) == 1 + input_tensor = input_tensor[0] + self._output_features = input_tensor + + self.classifier.register_forward_hook(hook_fn) + + def forward(self, x: torch.Tensor): + preds = self.base_module(x) + + if self._output_features is None: + raise RuntimeError("Target module did not receive any input during forward pass") + features, self._output_features = self._output_features, None + assert all([x == 1 for x in features.shape[2:]]), f'Invalid feature shape: {features.shape!r}' + features = torch.flatten(features, start_dim=1) + + return features, preds + + +def get_model(model: nn.Module, dummy_input: torch.Tensor): + assert isinstance(model, nn.Sequential) + head = model[-1] + wrapped_model = ModuleWrapper(model, head) + + logging.info(f'Input size: {dummy_input.shape!r}') + with torch.no_grad(): + dummy_embedding, dummy_preds = wrapped_model(dummy_input) + logging.info(f'Embedding size: {dummy_embedding.shape!r}') + logging.info(f'Preds size: {dummy_preds.shape!r}') + + return wrapped_model, (dummy_embedding, dummy_preds) + # print(model[-1]) From 77db215ee953725240783d2744c1ef849604c92d Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Mon, 8 Sep 2025 17:36:47 +0800 Subject: [PATCH 03/32] dev(narugo): fix convertion bug, ci skip --- zoo/pixai_tagger/export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 31c13a33d93..09f40c15212 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -75,7 +75,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): df_th.to_csv(os.path.join(upload_dir, 'thresholds.csv'), index=False) dummy_image = load_image(get_testfile('6125785.jpg'), mode='RGB', force_background='white') - dummy_input = handler.transform(dummy_image).unsqueeze(0) + dummy_input = handler.transform(dummy_image).unsqueeze(0).to(handler.device) wrapped_model, (conv_features, _) = get_model(handler.model, dummy_input) onnx_filename = os.path.join(upload_dir, 'model.onnx') with TemporaryDirectory() as td: From 02fb5b1f0202de7761bfd425e8bed1e9bbf69a8b Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Mon, 8 Sep 2025 17:42:18 +0800 Subject: [PATCH 04/32] dev(narugo): fix convertion bug, ci skip --- zoo/pixai_tagger/export.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 09f40c15212..70072916d1d 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -8,7 +8,7 @@ import torch from ditk import logging from hbutils.system import TemporaryDirectory -from hfutils.operate import get_hf_client +from hfutils.operate import get_hf_client, upload_directory_as_directory from imgutils.data import load_image from imgutils.preprocess import parse_torchvision_transforms @@ -77,6 +77,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): dummy_image = load_image(get_testfile('6125785.jpg'), mode='RGB', force_background='white') dummy_input = handler.transform(dummy_image).unsqueeze(0).to(handler.device) wrapped_model, (conv_features, _) = get_model(handler.model, dummy_input) + conv_features = conv_features.detach().cpu() onnx_filename = os.path.join(upload_dir, 'model.onnx') with TemporaryDirectory() as td: temp_model_onnx = os.path.join(td, 'model.onnx') @@ -111,7 +112,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): onnx.save(model, onnx_filename) session = onnxruntime.InferenceSession(onnx_filename) - o_embeddings, = session.run(['embedding'], {'input': dummy_input.numpy()}) + o_embeddings, = session.run(['embedding'], {'input': dummy_input.detach().cpu().numpy()}) emb_1 = o_embeddings / np.linalg.norm(o_embeddings, axis=-1, keepdims=True) emb_2 = conv_features.numpy() / np.linalg.norm(conv_features.numpy(), axis=-1, keepdims=True) emb_sims = (emb_1 * emb_2).sum() @@ -119,7 +120,14 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): assert emb_sims >= 0.98, f'Similarity of the embeddings is {emb_sims:.5f}, ONNX validation failed.' os.system(f'tree {upload_dir!r}') - input() + upload_directory_as_directory( + repo_id=dst_repo, + repo_type='model', + local_directory=upload_dir, + path_in_repo='.', + message=f'Upload ONNX export of model {src_repo!r}' + ) + # print(df_tags) # pprint(handler.index_to_tag_map) From 9fe4c391e37f0fcbec44047ec796f52d3e3ec443 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 00:47:18 +0800 Subject: [PATCH 05/32] dev(narugo): add a rough README --- zoo/pixai_tagger/export.py | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 70072916d1d..97f0d362edd 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -9,6 +9,7 @@ from ditk import logging from hbutils.system import TemporaryDirectory from hfutils.operate import get_hf_client, upload_directory_as_directory +from hfutils.repository import hf_hub_repo_url from imgutils.data import load_image from imgutils.preprocess import parse_torchvision_transforms @@ -119,7 +120,35 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): logging.info(f'Similarity of the embeddings is {emb_sims:.5f}.') assert emb_sims >= 0.98, f'Similarity of the embeddings is {emb_sims:.5f}, ONNX validation failed.' - os.system(f'tree {upload_dir!r}') + with open(os.path.join(upload_dir, 'README.md'), 'w') as f: + print('---', file=f) + print('pipeline_tag: image-classification', file=f) + print('base_model:', file=f) + print(f'- {src_repo}', file=f) + print('language:', file=f) + print('- en', file=f) + print('tags:', file=f) + print('- image', file=f) + print('- dghs-imgutils', file=f) + print('library_name: dghs-imgutils', file=f) + print('license: mit', file=f) + print('---', file=f) + print('', file=f) + + print(f'PixAI-Tagger ONNX Version for {src_repo}', file=f) + print(f'', file=f) + + print(f'This is the ONNX-exported version of PixAI\'s tagger ' + f'[{src_repo}]({hf_hub_repo_url(repo_id=src_repo, repo_type="model")}).', file=f) + print(f'', file=f) + + print(f'# How To Use', file=f) + print(f'', file=f) + print(f'```shell', file=f) + print(f'pip install -U dghs-imgutils', file=f) + print(f'```', file=f) + print(f'', file=f) + upload_directory_as_directory( repo_id=dst_repo, repo_type='model', From 0ff18ddba0a5ff4fc293b5eb72a7423c7f764b66 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 00:47:24 +0800 Subject: [PATCH 06/32] dev(narugo): add a rough README, ci skip --- zoo/camie/export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zoo/camie/export.py b/zoo/camie/export.py index 27657f379ec..884d20c34f5 100644 --- a/zoo/camie/export.py +++ b/zoo/camie/export.py @@ -19,7 +19,7 @@ from natsort import natsorted from safetensors.torch import save_model from thop import profile, clever_format - + from imgutils.data import load_image from imgutils.preprocess import parse_pillow_transforms, create_torchvision_transforms, parse_torchvision_transforms from imgutils.preprocess.pillow import PillowPadToSize, PillowToTensor, PillowCompose From 4e4e0ab450feb0204cb87d2afaec5fb51616b96c Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 00:47:26 +0800 Subject: [PATCH 07/32] dev(narugo): add a rough README, ci skip --- zoo/camie/export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zoo/camie/export.py b/zoo/camie/export.py index 884d20c34f5..27657f379ec 100644 --- a/zoo/camie/export.py +++ b/zoo/camie/export.py @@ -19,7 +19,7 @@ from natsort import natsorted from safetensors.torch import save_model from thop import profile, clever_format - + from imgutils.data import load_image from imgutils.preprocess import parse_pillow_transforms, create_torchvision_transforms, parse_torchvision_transforms from imgutils.preprocess.pillow import PillowPadToSize, PillowToTensor, PillowCompose From ff0f6dc7d8d6bbada9fc06296a599839dd279062 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 01:01:27 +0800 Subject: [PATCH 08/32] dev(narugo): add more profiles, ci skip --- requirements-zoo.txt | 3 +- zoo/pixai_tagger/export.py | 74 +++++++++++++++++++++++++++++++------- zoo/utils/__init__.py | 1 + zoo/utils/profile.py | 39 ++++++++++++++++++++ 4 files changed, 103 insertions(+), 14 deletions(-) create mode 100644 zoo/utils/profile.py diff --git a/requirements-zoo.txt b/requirements-zoo.txt index 515110506c7..fa598a87966 100644 --- a/requirements-zoo.txt +++ b/requirements-zoo.txt @@ -27,4 +27,5 @@ tabulate git+https://github.com/deepghs/waifuc.git@main#egg=waifuc pyquery httpx -onnxslim==0.1.32 \ No newline at end of file +onnxslim==0.1.32 +calflops \ No newline at end of file diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 97f0d362edd..f8dfd2f7988 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -7,14 +7,19 @@ import pandas as pd import torch from ditk import logging +from hbutils.string import titleize from hbutils.system import TemporaryDirectory +from hbutils.testing import vpip from hfutils.operate import get_hf_client, upload_directory_as_directory -from hfutils.repository import hf_hub_repo_url +from hfutils.repository import hf_hub_repo_url, hf_hub_repo_file_url +from hfutils.utils import hf_normpath +from huggingface_hub import hf_hub_url +from thop import clever_format from imgutils.data import load_image from imgutils.preprocess import parse_torchvision_transforms from zoo.pixai_tagger.tags import load_tags -from zoo.utils import onnx_optimize, get_testfile +from zoo.utils import onnx_optimize, get_testfile, torch_model_profile_via_calflops from .min_script import EndpointHandler from .onnx import get_model @@ -25,7 +30,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): hf_client.create_repo(repo_id=dst_repo, repo_type='model', private=True) handler = EndpointHandler(repo_id=src_repo) - + meta_info = {} with TemporaryDirectory() as upload_dir: preprocessor = handler.transform preprocessor_file = os.path.join(upload_dir, 'preprocessor.json') @@ -58,16 +63,16 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): logging.info(f'Tags:\n{df_tags}') df_tags.to_csv(os.path.join(upload_dir, 'selected_tags.csv'), index=False) + d_category_names = { + 0: 'general', + 4: 'character', + } with open(os.path.join(upload_dir, 'categories.json'), 'w') as f: json.dump([ { - "category": 0, - "name": "general" - }, - { - "category": 4, - "name": "character" - }, + "category": cate_id, + "name": cate_name, + } for cate_id, cate_name in sorted(d_category_names.items()) ], f, sort_keys=True, ensure_ascii=False, indent=4) df_th = pd.DataFrame([ {'category': 0, 'name': 'general', 'threshold': handler.default_general_threshold}, @@ -77,6 +82,15 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): dummy_image = load_image(get_testfile('6125785.jpg'), mode='RGB', force_background='white') dummy_input = handler.transform(dummy_image).unsqueeze(0).to(handler.device) + flops, params, macs = torch_model_profile_via_calflops(model=handler.model, input_=dummy_input) + meta_info['flops'] = flops + meta_info['params'] = params + meta_info['macs'] = macs + new_meta_file = os.path.join(upload_dir, 'meta.json') + logging.info(f'Saving metadata to {new_meta_file!r} ...') + with open(new_meta_file, 'w') as f: + json.dump(meta_info, f, indent=4, sort_keys=True, ensure_ascii=False) + wrapped_model, (conv_features, _) = get_model(handler.model, dummy_input) conv_features = conv_features.detach().cpu() onnx_filename = os.path.join(upload_dir, 'model.onnx') @@ -135,17 +149,51 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): print('---', file=f) print('', file=f) - print(f'PixAI-Tagger ONNX Version for {src_repo}', file=f) + print(f'# ONNX Version for {src_repo}', file=f) print(f'', file=f) print(f'This is the ONNX-exported version of PixAI\'s tagger ' f'[{src_repo}]({hf_hub_repo_url(repo_id=src_repo, repo_type="model")}).', file=f) print(f'', file=f) - print(f'# How To Use', file=f) + s_flops, s_params, s_macs = clever_format([flops, params, macs], "%.1f") + print(f'## Model Details', file=f) + print(f'', file=f) + print(f'- **Model Type:** Multilabel Image classification / feature backbone', file=f) + print(f'- **Model Stats:**', file=f) + print(f' - Params: {s_params}', file=f) + print(f' - FLOPs / MACs: {s_flops} / {s_macs}', file=f) + print(f' - Image size: {dummy_input.shape[-1]} x {dummy_input.shape[-2]}', file=f) + print(f' - Tags Count: {len(df_tags)}', file=f) + for category in sorted(set(df_tags['category'])): + print(f' - {titleize(d_category_names[category])} (#{category}) Tags Count: ' + f'{len(df_tags[df_tags["category"] == category])}', file=f) + print(f'', file=f) + + print(f'## How to Use', file=f) + print(f'', file=f) + + imgutils_version = str(vpip('dghs-imgutils')._actual_version) + sample_input = dummy_image + if min(sample_input.width, sample_input.height) > 640: + r = min(sample_input.width, sample_input.height) / 640 + new_width = int(sample_input.width / r) + new_height = int(sample_input.height / r) + sample_input = sample_input.resize((new_width, new_height)) + sample_input_file = os.path.join(upload_dir, 'sample.webp') + sample_input_relfile = hf_normpath(os.path.relpath(sample_input_file, upload_dir)) + sample_input.save(sample_input_file) + sample_input_url = hf_hub_url(repo_id=dst_repo, repo_type='model', filename=sample_input_relfile) + sample_input_page_url = hf_hub_repo_file_url(repo_id=dst_repo, repo_type='model', path=sample_input_relfile) + + print(f'We provided a sample image for our code samples, ' + f'you can find it [here]({sample_input_page_url}).', file=f) + print(f'', file=f) + + print(f'Install [dghs-imgutils](https://github.com/deepghs/imgutils) with the following command', file=f) print(f'', file=f) print(f'```shell', file=f) - print(f'pip install -U dghs-imgutils', file=f) + print(f'pip install \'dghs-imgutils>={imgutils_version}\' torch huggingface_hub timm pillow pandas', file=f) print(f'```', file=f) print(f'', file=f) diff --git a/zoo/utils/__init__.py b/zoo/utils/__init__.py index d699c746f2d..40d974e41f0 100644 --- a/zoo/utils/__init__.py +++ b/zoo/utils/__init__.py @@ -2,4 +2,5 @@ from .lr import get_init_lr, get_dynamic_lr_scheduler, LRTyping from .onnx import onnx_quick_export from .optimize import onnx_optimize +from .profile import torch_model_profile_via_thop, torch_model_profile_via_calflops from .testfile import get_testfile diff --git a/zoo/utils/profile.py b/zoo/utils/profile.py new file mode 100644 index 00000000000..169aabff244 --- /dev/null +++ b/zoo/utils/profile.py @@ -0,0 +1,39 @@ +import torch +from ditk import logging +from thop import profile, clever_format + + +def torch_model_profile_via_thop(model, input_): + with torch.no_grad(): + flops, params = profile(model, (input_,)) + + s_flops, s_params = clever_format([flops, params], "%.1f") + logging.info(f'Params: {s_params}, FLOPs: {s_flops}.') + + return flops, params + + +def torch_model_profile_via_calflops(model, input_): + from calflops import calculate_flops + flops, macs, params = calculate_flops( + model=model, + input_shape=tuple(input_.shape), + output_as_string=False, + print_detailed=False, + # output_as_string=True, + # output_precision=4 + ) + s_flops, s_params, s_macs = clever_format([flops, params, macs], "%.1f") + logging.info(f'Params: {s_params}, FLOPs: {s_flops}, MACs: {s_macs}.') + return flops, params, macs + + +if __name__ == '__main__': + logging.try_init_root(level=logging.INFO) + from timm import create_model + + # model = create_model('hf-hub:animetimm/swinv2_base_window8_256.dbv4-full', pretrained=False) + model = create_model('caformer_b36.sail_in22k_ft_in1k_384', pretrained=False) + dummy_input = torch.randn(1, 3, 448, 448) + print(torch_model_profile_via_thop(model, dummy_input)) + print(torch_model_profile_via_calflops(model, dummy_input)) From 22eb21db56548d25360826a43896b380bc7629e6 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 01:04:04 +0800 Subject: [PATCH 09/32] dev(narugo): add more profiles, ci skip --- requirements-zoo.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-zoo.txt b/requirements-zoo.txt index fa598a87966..cb8633904fa 100644 --- a/requirements-zoo.txt +++ b/requirements-zoo.txt @@ -13,7 +13,7 @@ tensorboard einops thop accelerate -timm~=0.6.13 +timm>=1 ftfy regex git+https://github.com/openai/CLIP.git From 82cf1103d9878b721dcf53fd34a3017f1dc64ede Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 01:06:20 +0800 Subject: [PATCH 10/32] dev(narugo): add more profiles, ci skip --- requirements-zoo.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements-zoo.txt b/requirements-zoo.txt index cb8633904fa..48dd4ff5acd 100644 --- a/requirements-zoo.txt +++ b/requirements-zoo.txt @@ -28,4 +28,5 @@ git+https://github.com/deepghs/waifuc.git@main#egg=waifuc pyquery httpx onnxslim==0.1.32 -calflops \ No newline at end of file +calflops +transformers \ No newline at end of file From 18a9bad5a9070d29d540636fd9bbb57a64f9ac4a Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 01:20:05 +0800 Subject: [PATCH 11/32] dev(narugo): add more profiles, ci skip --- zoo/pixai_tagger/export.py | 7 ++++--- zoo/pixai_tagger/min_script.py | 3 ++- zoo/pixai_tagger/onnx.py | 29 +++++++++++++++++++++-------- 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index f8dfd2f7988..15e5f7f8eb2 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -91,7 +91,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): with open(new_meta_file, 'w') as f: json.dump(meta_info, f, indent=4, sort_keys=True, ensure_ascii=False) - wrapped_model, (conv_features, _) = get_model(handler.model, dummy_input) + wrapped_model, (conv_features, _, _) = get_model(handler.model, dummy_input) conv_features = conv_features.detach().cpu() onnx_filename = os.path.join(upload_dir, 'model.onnx') with TemporaryDirectory() as td: @@ -102,11 +102,12 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): dummy_input, temp_model_onnx, input_names=['input'], - output_names=['embedding', 'output'], + output_names=['embedding', 'logits', 'prediction'], dynamic_axes={ 'input': {0: 'batch_size'}, 'embedding': {0: 'batch_size'}, - 'output': {0: 'batch_size'}, + 'logits': {0: 'batch_size'}, + 'prediction': {0: 'batch_size'}, }, opset_version=14, do_constant_folding=True, diff --git a/zoo/pixai_tagger/min_script.py b/zoo/pixai_tagger/min_script.py index 40cf338178d..1c06ace1fcf 100644 --- a/zoo/pixai_tagger/min_script.py +++ b/zoo/pixai_tagger/min_script.py @@ -22,10 +22,11 @@ def __init__(self, input_dim, num_classes): self.input_dim = input_dim self.num_classes = num_classes self.head = torch.nn.Sequential(torch.nn.Linear(input_dim, num_classes)) + self.sigmoid = torch.nn.Sigmoid() def forward(self, x): logits = self.head(x) - probs = torch.nn.functional.sigmoid(logits) + probs = self.sigmoid(logits) return probs diff --git a/zoo/pixai_tagger/onnx.py b/zoo/pixai_tagger/onnx.py index 759b9b7aaf7..1644322ac8f 100644 --- a/zoo/pixai_tagger/onnx.py +++ b/zoo/pixai_tagger/onnx.py @@ -5,44 +5,57 @@ class ModuleWrapper(nn.Module): - def __init__(self, base_module: nn.Module, classifier: nn.Module): + def __init__(self, base_module: nn.Module, classifier: nn.Module, sigmoid: nn.Module): super().__init__() self.base_module = base_module self.classifier = classifier + self.sigmoid = sigmoid self._output_features = None + self._output_logits = None self._register_hook() def _register_hook(self): - def hook_fn(module, input_tensor, output_tensor): + def hook_fn_embeddings(module, input_tensor, output_tensor): assert isinstance(input_tensor, tuple) and len(input_tensor) == 1 input_tensor = input_tensor[0] self._output_features = input_tensor - self.classifier.register_forward_hook(hook_fn) + self.classifier.register_forward_hook(hook_fn_embeddings) + + def hook_fn_logits(module, input_tensor, output_tensor): + assert isinstance(input_tensor, tuple) and len(input_tensor) == 1 + input_tensor = input_tensor[0] + self._output_logits = input_tensor + + self.sigmoid.register_forward_hook(hook_fn_logits) def forward(self, x: torch.Tensor): preds = self.base_module(x) if self._output_features is None: - raise RuntimeError("Target module did not receive any input during forward pass") + raise RuntimeError("Target module did not receive any input during forward pass (features)") + if self._output_logits is None: + raise RuntimeError("Target module did not receive any input during forward pass (logits)") features, self._output_features = self._output_features, None + logits, self._output_logits = self._output_logits, None assert all([x == 1 for x in features.shape[2:]]), f'Invalid feature shape: {features.shape!r}' features = torch.flatten(features, start_dim=1) - return features, preds + return features, logits, preds def get_model(model: nn.Module, dummy_input: torch.Tensor): assert isinstance(model, nn.Sequential) head = model[-1] - wrapped_model = ModuleWrapper(model, head) + wrapped_model = ModuleWrapper(model, head, head.sigmoid) logging.info(f'Input size: {dummy_input.shape!r}') with torch.no_grad(): - dummy_embedding, dummy_preds = wrapped_model(dummy_input) + dummy_embedding, dummy_logits, dummy_preds = wrapped_model(dummy_input) logging.info(f'Embedding size: {dummy_embedding.shape!r}') + logging.info(f'Logits size: {dummy_preds.shape!r}') logging.info(f'Preds size: {dummy_preds.shape!r}') - return wrapped_model, (dummy_embedding, dummy_preds) + return wrapped_model, (dummy_embedding, dummy_logits, dummy_preds) # print(model[-1]) From 150b42f8d6a4821ba60fe532113d930e44e9da77 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 09:51:24 +0800 Subject: [PATCH 12/32] dev(narugo): add some more meta information --- zoo/pixai_tagger/export.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 15e5f7f8eb2..322b1dd22bf 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -31,6 +31,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): handler = EndpointHandler(repo_id=src_repo) meta_info = {} + meta_info['repo_id'] = src_repo with TemporaryDirectory() as upload_dir: preprocessor = handler.transform preprocessor_file = os.path.join(upload_dir, 'preprocessor.json') @@ -82,6 +83,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): dummy_image = load_image(get_testfile('6125785.jpg'), mode='RGB', force_background='white') dummy_input = handler.transform(dummy_image).unsqueeze(0).to(handler.device) + meta_info['input_size'] = dummy_input.shape[-1] flops, params, macs = torch_model_profile_via_calflops(model=handler.model, input_=dummy_input) meta_info['flops'] = flops meta_info['params'] = params @@ -91,8 +93,12 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): with open(new_meta_file, 'w') as f: json.dump(meta_info, f, indent=4, sort_keys=True, ensure_ascii=False) - wrapped_model, (conv_features, _, _) = get_model(handler.model, dummy_input) + wrapped_model, (conv_features, conv_logits, conv_preds) = get_model(handler.model, dummy_input) conv_features = conv_features.detach().cpu() + conv_logits = conv_logits.detach().cpu() + conv_preds = conv_preds.detach().cpu() + meta_info['num_features'] = conv_features.shape[-1] + meta_info['num_classes'] = conv_preds.shape[-1] onnx_filename = os.path.join(upload_dir, 'model.onnx') with TemporaryDirectory() as td: temp_model_onnx = os.path.join(td, 'model.onnx') From 9f266c991f8e5ccd5c27acfe5b00f192523c79cf Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 09:53:05 +0800 Subject: [PATCH 13/32] dev(narugo): add some more meta information, ci skip --- zoo/pixai_tagger/export.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 322b1dd22bf..487346e360f 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -68,6 +68,10 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): 0: 'general', 4: 'character', } + d_category_thresholds = { + 0: handler.default_general_threshold, + 4: handler.default_character_threshold, + } with open(os.path.join(upload_dir, 'categories.json'), 'w') as f: json.dump([ { @@ -76,8 +80,11 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): } for cate_id, cate_name in sorted(d_category_names.items()) ], f, sort_keys=True, ensure_ascii=False, indent=4) df_th = pd.DataFrame([ - {'category': 0, 'name': 'general', 'threshold': handler.default_general_threshold}, - {'category': 4, 'name': 'character', 'threshold': handler.default_character_threshold}, + { + "category": cate_id, + "name": cate_name, + 'threshold': d_category_thresholds[cate_id] + } for cate_id, cate_name in sorted(d_category_names.items()) ]) df_th.to_csv(os.path.join(upload_dir, 'thresholds.csv'), index=False) From 93bce42a7bd37e52779bc68d37b53dc5caa91086 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 09:59:13 +0800 Subject: [PATCH 14/32] dev(narugo): add some more meta information, ci skip --- zoo/pixai_tagger/export.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 487346e360f..9c9dfe6396b 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -95,10 +95,6 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): meta_info['flops'] = flops meta_info['params'] = params meta_info['macs'] = macs - new_meta_file = os.path.join(upload_dir, 'meta.json') - logging.info(f'Saving metadata to {new_meta_file!r} ...') - with open(new_meta_file, 'w') as f: - json.dump(meta_info, f, indent=4, sort_keys=True, ensure_ascii=False) wrapped_model, (conv_features, conv_logits, conv_preds) = get_model(handler.model, dummy_input) conv_features = conv_features.detach().cpu() @@ -106,6 +102,11 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): conv_preds = conv_preds.detach().cpu() meta_info['num_features'] = conv_features.shape[-1] meta_info['num_classes'] = conv_preds.shape[-1] + new_meta_file = os.path.join(upload_dir, 'meta.json') + logging.info(f'Saving metadata to {new_meta_file!r} ...') + with open(new_meta_file, 'w') as f: + json.dump(meta_info, f, indent=4, sort_keys=True, ensure_ascii=False) + onnx_filename = os.path.join(upload_dir, 'model.onnx') with TemporaryDirectory() as td: temp_model_onnx = os.path.join(td, 'model.onnx') From 6c85e626b28a9fd89cbef37fdeac3f242f5f996c Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 10:36:15 +0800 Subject: [PATCH 15/32] dev(narugo): try add inference code, ci skip --- imgutils/generic/multilabel_timm.py | 3 +- imgutils/tagging/__init__.py | 1 + imgutils/tagging/pixai.py | 183 ++++++++++++++++++++++++++++ zoo/pixai_tagger/export.py | 5 +- 4 files changed, 189 insertions(+), 3 deletions(-) create mode 100644 imgutils/tagging/pixai.py diff --git a/imgutils/generic/multilabel_timm.py b/imgutils/generic/multilabel_timm.py index 294394a0714..b1e3a3e7623 100644 --- a/imgutils/generic/multilabel_timm.py +++ b/imgutils/generic/multilabel_timm.py @@ -29,6 +29,7 @@ from typing import Optional, Literal, Dict, Any, Union import pandas as pd +from hbutils.design import SingletonMark from hbutils.string import titleize from hfutils.repository import hf_hub_repo_url from huggingface_hub import hf_hub_download @@ -64,7 +65,7 @@ def _check_gradio_env(): f'Please install it with `pip install dghs-imgutils[demo]`.') -FMT_UNSET = object() +FMT_UNSET = SingletonMark('FMT_UNSET') class MultiLabelTIMMModel: diff --git a/imgutils/tagging/__init__.py b/imgutils/tagging/__init__.py index 289dc6775e2..f4124f1c41e 100644 --- a/imgutils/tagging/__init__.py +++ b/imgutils/tagging/__init__.py @@ -18,4 +18,5 @@ from .mldanbooru import get_mldanbooru_tags from .order import sort_tags from .overlap import drop_overlap_tags +from .pixai import get_pixai_tags from .wd14 import get_wd14_tags, convert_wd14_emb_to_prediction, denormalize_wd14_emb diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py new file mode 100644 index 00000000000..3004d8e45ee --- /dev/null +++ b/imgutils/tagging/pixai.py @@ -0,0 +1,183 @@ +import json +from typing import Union, Dict, Any + +import pandas as pd +from hbutils.design import SingletonMark +from huggingface_hub import hf_hub_download +from huggingface_hub.errors import EntryNotFoundError + +from imgutils.data import ImageTyping, load_image +from imgutils.preprocess import create_pillow_transforms +from imgutils.utils import open_onnx_model, ts_lru_cache, vreplace + +FMT_UNSET = SingletonMark('FMT_UNSET') + + +def _get_repo_id(model_name: str) -> str: + if '/' in model_name: + return model_name + else: + return f'deepghs/pixai-tagger-{model_name}-onnx' + + +@ts_lru_cache() +def _open_onnx_model(model_name: str): + """ + Load the ONNX model from Hugging Face Hub. + + :return: The loaded ONNX model + :rtype: object + """ + return open_onnx_model(hf_hub_download( + repo_id=_get_repo_id(model_name), + repo_type='model', + filename='model.onnx', + )) + + +@ts_lru_cache() +def _open_tags(model_name: str) -> pd.DataFrame: + return pd.read_csv(hf_hub_download( + repo_id=_get_repo_id(model_name), + repo_type='model', + filename='selected_tags.csv', + )) + + +@ts_lru_cache() +def _open_preprocess(model_name: str): + with open(hf_hub_download( + repo_id=_get_repo_id(model_name), + repo_type='model', + filename='preprocess.json' + ), 'r') as f: + data_ = json.load(f) + return create_pillow_transforms(data_['stages']) + + +@ts_lru_cache() +def _open_default_category_thresholds(model_name: str) -> Union[Dict[int, float], Dict[int, str]]: + """ + Load default category thresholds from the Hugging Face Hub. + + :return: Dictionary mapping category IDs to threshold values + :rtype: dict + """ + try: + df_category_thresholds = pd.read_csv(hf_hub_download( + repo_id=_get_repo_id(model_name), + repo_type='model', + filename='thresholds.csv' + ), keep_default_na=False) + except (EntryNotFoundError,): + _default_category_thresholds = {} + _category_names = {} + else: + _default_category_thresholds = {} + _category_names = {} + for item in df_category_thresholds.to_dict('records'): + if item['category'] not in _default_category_thresholds: + _default_category_thresholds[item['category']] = item['threshold'] + _category_names[item['category']] = item['name'] + + return _default_category_thresholds, _category_names + + +def _raw_predict(image: ImageTyping, model_name: str): + """ + Make a raw prediction with the model. + + :param image: The input image + :type image: ImageTyping + :param preprocessor: Which preprocessor to use ('test' or 'val') + :type preprocessor: Literal['test', 'val'] + + :return: Dictionary of model outputs + :rtype: dict + :raises ValueError: If an unknown preprocessor is specified + """ + image = load_image(image, force_background='white', mode='RGB') + model = _open_onnx_model(model_name=model_name) + trans = _open_preprocess(model_name=model_name) + input_ = trans(image)[None, ...] + output_names = [output.name for output in model.get_outputs()] + output_values = model.run(output_names, {'input': input_}) + return {name: value[0] for name, value in zip(output_names, output_values)} + + +def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', + thresholds: Union[float, Dict[Any, float]] = None, fmt=FMT_UNSET): + """ + Make a prediction and format the results. + + This method processes an image through the model and applies thresholds to determine + which tags to include in the results. The output format can be customized using the fmt parameter. + + :param image: The input image + :type image: ImageTyping + :param preprocessor: Which preprocessor to use ('test' or 'val') + :type preprocessor: Literal['test', 'val'] + :param thresholds: Threshold values for tag confidence. Can be a single float applied to all categories + or a dictionary mapping category IDs or names to threshold values + :type thresholds: Union[float, Dict[Any, float]] + :param use_tag_thresholds: Whether to use tag-level thresholds if available + :type use_tag_thresholds: bool + :param fmt: Output format specification. Can be a tuple of category names to include, + or FMT_UNSET to use all categories + :type fmt: Any + + :return: Formatted prediction results according to the fmt parameter + :rtype: Any + + .. note:: + The fmt argument can include the following keys: + + - Category names: dicts containing category-specific tags and their confidences + - ``tag``: a dict containing all tags across categories and their confidences + - ``embedding``: a 1-dim embedding of image, recommended for index building after L2 normalization + - ``logits``: a 1-dim logits result of image + - ``prediction``: a 1-dim prediction result of image + + You can extract specific category predictions or all tags based on your needs. + + For more details see documentation of :func:`multilabel_timm_predict`. + """ + df_tags = _open_tags(model_name=model_name) + values = _raw_predict(image, model_name=model_name) + prediction = values['prediction'] + tags = {} + + default_category_thresholds, category_names = _open_default_category_thresholds() + if fmt is FMT_UNSET: + fmt = tuple(category_names[category] for category in sorted(set(df_tags['category'].tolist()))) + + for category in sorted(set(df_tags['category'].tolist())): + mask = df_tags['category'] == category + tag_names = df_tags['name'][mask] + category_pred = prediction[mask] + + if isinstance(thresholds, float): + category_threshold = thresholds + elif isinstance(thresholds, dict) and \ + (category in thresholds or category_names[category] in thresholds): + if category in thresholds: + category_threshold = thresholds[category] + elif category_names[category] in thresholds: + category_threshold = thresholds[category_names[category]] + else: + assert False, 'Should not reach this line' # pragma: no cover + else: + if category in default_category_thresholds: + category_threshold = default_category_thresholds[category] + else: + category_threshold = 0.4 + + mask = category_pred >= category_threshold + tag_names = tag_names[mask].tolist() + category_pred = category_pred[mask].tolist() + cate_tags = dict(sorted(zip(tag_names, category_pred), key=lambda x: (-x[1], x[0]))) + values[category_names[category]] = cate_tags + tags.update(cate_tags) + + values['tag'] = tags + return vreplace(fmt, values) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 9c9dfe6396b..f16ed4ef6f0 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -34,7 +34,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): meta_info['repo_id'] = src_repo with TemporaryDirectory() as upload_dir: preprocessor = handler.transform - preprocessor_file = os.path.join(upload_dir, 'preprocessor.json') + preprocessor_file = os.path.join(upload_dir, 'preprocess.json') logging.info(f'Dumping preprocessor:\n{preprocessor}\nto file {preprocessor_file!r}.') with open(preprocessor_file, 'w') as f: json.dump({ @@ -217,7 +217,8 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): repo_type='model', local_directory=upload_dir, path_in_repo='.', - message=f'Upload ONNX export of model {src_repo!r}' + message=f'Upload ONNX export of model {src_repo!r}', + clear=True, ) # print(df_tags) From 313a6536ae6b35ce2b92585a35cc906bfc463c1d Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 10:46:44 +0800 Subject: [PATCH 16/32] dev(narugo): fix known bugs --- imgutils/tagging/pixai.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py index 3004d8e45ee..9e1f564c884 100644 --- a/imgutils/tagging/pixai.py +++ b/imgutils/tagging/pixai.py @@ -63,6 +63,8 @@ def _open_default_category_thresholds(model_name: str) -> Union[Dict[int, float] :return: Dictionary mapping category IDs to threshold values :rtype: dict """ + _default_category_thresholds: Dict[int, float] = {} + _category_names: Dict[int, str] = {} try: df_category_thresholds = pd.read_csv(hf_hub_download( repo_id=_get_repo_id(model_name), @@ -70,11 +72,8 @@ def _open_default_category_thresholds(model_name: str) -> Union[Dict[int, float] filename='thresholds.csv' ), keep_default_na=False) except (EntryNotFoundError,): - _default_category_thresholds = {} - _category_names = {} + pass else: - _default_category_thresholds = {} - _category_names = {} for item in df_category_thresholds.to_dict('records'): if item['category'] not in _default_category_thresholds: _default_category_thresholds[item['category']] = item['threshold'] @@ -147,7 +146,7 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', prediction = values['prediction'] tags = {} - default_category_thresholds, category_names = _open_default_category_thresholds() + default_category_thresholds, category_names = _open_default_category_thresholds(model_name=model_name) if fmt is FMT_UNSET: fmt = tuple(category_names[category] for category in sorted(set(df_tags['category'].tolist()))) From db8a1bae2ea418a436d415aa57cf9d5ae1923854 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 10:56:53 +0800 Subject: [PATCH 17/32] dev(narugo): add pixai docs --- docs/source/api_doc/tagging/index.rst | 1 + docs/source/api_doc/tagging/pixai.rst | 15 +++ imgutils/tagging/pixai.py | 187 ++++++++++++++++++++------ 3 files changed, 164 insertions(+), 39 deletions(-) create mode 100644 docs/source/api_doc/tagging/pixai.rst diff --git a/docs/source/api_doc/tagging/index.rst b/docs/source/api_doc/tagging/index.rst index 34c51a733c8..b6213451a3f 100644 --- a/docs/source/api_doc/tagging/index.rst +++ b/docs/source/api_doc/tagging/index.rst @@ -12,6 +12,7 @@ imgutils.tagging mldanbooru wd14 camie + pixai deepdanbooru deepgelbooru format diff --git a/docs/source/api_doc/tagging/pixai.rst b/docs/source/api_doc/tagging/pixai.rst new file mode 100644 index 00000000000..b87c782d432 --- /dev/null +++ b/docs/source/api_doc/tagging/pixai.rst @@ -0,0 +1,15 @@ +imgutils.tagging.pixai +==================================== + +.. currentmodule:: imgutils.tagging.pixai + +.. automodule:: imgutils.tagging.pixai + + +get_pixai_tags +---------------------- + +.. autofunction:: get_pixai_tags + + + diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py index 9e1f564c884..85ef5e2cc34 100644 --- a/imgutils/tagging/pixai.py +++ b/imgutils/tagging/pixai.py @@ -1,5 +1,29 @@ +""" +Overview: + This module provides utilities for image tagging using PixAI taggers, which are specialized models + for analyzing anime-style images and extracting relevant tags. The module supports loading ONNX + models from Hugging Face Hub and processing images to generate categorized tags with confidence scores. + + The models are originally developed by the PixAI team and available at + `pixai-labs `_ on Hugging Face. This module uses ONNX-converted + versions of these models for efficient inference, available at + `deepghs `_ repositories. + + Example:: + >>> from imgutils.tagging.pixai import get_pixai_tags + >>> # Get tags with default thresholds + >>> result = get_pixai_tags('path/to/anime_image.jpg', model_name='v0.9') + >>> general_tags, character_tags = result + >>> print("General tags:", general_tags) + >>> print("Character tags:", character_tags) + + >>> # Get all tags in a single dictionary + >>> all_tags = get_pixai_tags('path/to/image.jpg', fmt='tag') + >>> print("All tags:", all_tags) +""" + import json -from typing import Union, Dict, Any +from typing import Union, Dict, Any, Tuple import pandas as pd from hbutils.design import SingletonMark @@ -14,6 +38,21 @@ def _get_repo_id(model_name: str) -> str: + """ + Get the repository ID for the specified model name. + + :param model_name: Name of the model (e.g., 'v0.9') or full repository path + :type model_name: str + + :return: Full repository ID for Hugging Face Hub + :rtype: str + + Example:: + >>> _get_repo_id('v0.9') + 'deepghs/pixai-tagger-v0.9-onnx' + >>> _get_repo_id('custom/model-repo') + 'custom/model-repo' + """ if '/' in model_name: return model_name else: @@ -23,10 +62,16 @@ def _get_repo_id(model_name: str) -> str: @ts_lru_cache() def _open_onnx_model(model_name: str): """ - Load the ONNX model from Hugging Face Hub. + Load the ONNX model from Hugging Face Hub with caching. + + This function downloads and loads the ONNX model file for the specified PixAI tagger. + Results are cached to avoid repeated downloads and model loading. - :return: The loaded ONNX model - :rtype: object + :param model_name: Name of the model to load + :type model_name: str + + :return: The loaded ONNX model session + :rtype: onnxruntime.InferenceSession """ return open_onnx_model(hf_hub_download( repo_id=_get_repo_id(model_name), @@ -37,6 +82,19 @@ def _open_onnx_model(model_name: str): @ts_lru_cache() def _open_tags(model_name: str) -> pd.DataFrame: + """ + Load the tag metadata from Hugging Face Hub with caching. + + This function downloads and loads the CSV file containing tag names and categories + for the specified model. The DataFrame contains columns for tag names, categories, + and other metadata. + + :param model_name: Name of the model + :type model_name: str + + :return: DataFrame containing tag information with columns like 'name', 'category' + :rtype: pd.DataFrame + """ return pd.read_csv(hf_hub_download( repo_id=_get_repo_id(model_name), repo_type='model', @@ -46,6 +104,17 @@ def _open_tags(model_name: str) -> pd.DataFrame: @ts_lru_cache() def _open_preprocess(model_name: str): + """ + Load the preprocessing pipeline configuration from Hugging Face Hub with caching. + + This function downloads the preprocessing configuration and creates a PIL transforms + pipeline for image preprocessing before model inference. + + :param model_name: Name of the model + :type model_name: str + + :return: Preprocessing transform pipeline + """ with open(hf_hub_download( repo_id=_get_repo_id(model_name), repo_type='model', @@ -56,12 +125,23 @@ def _open_preprocess(model_name: str): @ts_lru_cache() -def _open_default_category_thresholds(model_name: str) -> Union[Dict[int, float], Dict[int, str]]: +def _open_default_category_thresholds(model_name: str) -> Tuple[Dict[int, float], Dict[int, str]]: """ - Load default category thresholds from the Hugging Face Hub. + Load default category thresholds and names from the Hugging Face Hub with caching. - :return: Dictionary mapping category IDs to threshold values - :rtype: dict + This function attempts to load predefined threshold values for each category from + a CSV file. If the file doesn't exist, empty dictionaries are returned. + + :param model_name: Name of the model + :type model_name: str + + :return: Tuple containing (category_thresholds, category_names) dictionaries + :rtype: tuple[Dict[int, float], Dict[int, str]] + + Example:: + >>> thresholds, names = _open_default_category_thresholds('v0.9') + >>> print(thresholds) # {0: 0.35, 1: 0.4, ...} + >>> print(names) # {0: 'general', 1: 'character', ...} """ _default_category_thresholds: Dict[int, float] = {} _category_names: Dict[int, str] = {} @@ -84,16 +164,23 @@ def _open_default_category_thresholds(model_name: str) -> Union[Dict[int, float] def _raw_predict(image: ImageTyping, model_name: str): """ - Make a raw prediction with the model. + Make a raw prediction with the PixAI tagger model. + + This function preprocesses the input image and runs inference using the specified + ONNX model. It returns the raw model outputs without any post-processing or + threshold application. - :param image: The input image + :param image: The input image to analyze :type image: ImageTyping - :param preprocessor: Which preprocessor to use ('test' or 'val') - :type preprocessor: Literal['test', 'val'] + :param model_name: Name of the model to use for prediction + :type model_name: str - :return: Dictionary of model outputs + :return: Dictionary containing raw model outputs with keys like 'prediction', 'embedding', 'logits' :rtype: dict - :raises ValueError: If an unknown preprocessor is specified + + Example:: + >>> raw_output = _raw_predict('anime_image.jpg', 'v0.9') + >>> print(raw_output.keys()) # dict_keys(['prediction', 'embedding', 'logits']) """ image = load_image(image, force_background='white', mode='RGB') model = _open_onnx_model(model_name=model_name) @@ -107,39 +194,61 @@ def _raw_predict(image: ImageTyping, model_name: str): def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', thresholds: Union[float, Dict[Any, float]] = None, fmt=FMT_UNSET): """ - Make a prediction and format the results. + Extract tags from an image using PixAI tagger models. - This method processes an image through the model and applies thresholds to determine - which tags to include in the results. The output format can be customized using the fmt parameter. + This function processes an image through a PixAI tagger model and applies confidence + thresholds to determine which tags to include in the results. The output format can + be customized to return specific categories or all tags together. - :param image: The input image + :param image: The input image to analyze (file path, PIL Image, numpy array, etc.) :type image: ImageTyping - :param preprocessor: Which preprocessor to use ('test' or 'val') - :type preprocessor: Literal['test', 'val'] - :param thresholds: Threshold values for tag confidence. Can be a single float applied to all categories - or a dictionary mapping category IDs or names to threshold values - :type thresholds: Union[float, Dict[Any, float]] - :param use_tag_thresholds: Whether to use tag-level thresholds if available - :type use_tag_thresholds: bool - :param fmt: Output format specification. Can be a tuple of category names to include, - or FMT_UNSET to use all categories + :param model_name: Name or path of the PixAI tagger model to use + :type model_name: str + :param thresholds: Confidence threshold values. Can be a single float applied to all + categories, or a dictionary mapping category IDs/names to specific thresholds + :type thresholds: Union[float, Dict[Any, float]], optional + :param fmt: Output format specification. If FMT_UNSET, returns all available categories. + Can be a tuple of category names to include in output :type fmt: Any - :return: Formatted prediction results according to the fmt parameter + :return: Formatted prediction results. Default returns tuple of (general_tags, character_tags, ...) + based on available categories. Can return custom format based on fmt parameter :rtype: Any .. note:: - The fmt argument can include the following keys: - - - Category names: dicts containing category-specific tags and their confidences - - ``tag``: a dict containing all tags across categories and their confidences - - ``embedding``: a 1-dim embedding of image, recommended for index building after L2 normalization - - ``logits``: a 1-dim logits result of image - - ``prediction``: a 1-dim prediction result of image - - You can extract specific category predictions or all tags based on your needs. - - For more details see documentation of :func:`multilabel_timm_predict`. + The fmt parameter can include the following keys: + + - Category names (e.g., 'general', 'character'): dictionaries containing category-specific + tags and their confidence scores + - ``tag``: a dictionary containing all tags across categories and their confidences + - ``embedding``: a 1-dimensional embedding vector of the image, recommended for similarity + search after L2 normalization + - ``logits``: raw 1-dimensional logits output from the model + - ``prediction``: 1-dimensional prediction probabilities from the model + + Default category thresholds are used if not specified. These vary by model and category + but typically range from 0.35 to 0.5. + + Example:: + >>> from imgutils.tagging.pixai import get_pixai_tags + >>> + >>> # Get tags with default format (all categories) + >>> general_tags, character_tags = get_pixai_tags('anime_image.jpg', model_name='v0.9') + >>> print("General tags:", general_tags) + >>> print("Character tags:", character_tags) + >>> + >>> # Get all tags in a single dictionary + >>> all_tags = get_pixai_tags('image.jpg', fmt='tag') + >>> print("All tags:", all_tags) + >>> + >>> # Use custom thresholds + >>> result = get_pixai_tags('image.jpg', thresholds={'general': 0.3, 'character': 0.5}) + >>> + >>> # Get embedding for similarity search + >>> embedding = get_pixai_tags('image.jpg', fmt='embedding') + >>> # Normalize for cosine similarity + >>> import numpy as np + >>> normalized_embedding = embedding / np.linalg.norm(embedding) """ df_tags = _open_tags(model_name=model_name) values = _raw_predict(image, model_name=model_name) From c9ef10510e1c048696d3075214a2dbf61787c811 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 11:07:26 +0800 Subject: [PATCH 18/32] dev(narugo): add pixai docs, ci skip, fix some bugs --- imgutils/tagging/pixai.py | 25 +++++++++++++++++++++---- zoo/pixai_tagger/export.py | 2 +- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py index 85ef5e2cc34..8fada7289ab 100644 --- a/imgutils/tagging/pixai.py +++ b/imgutils/tagging/pixai.py @@ -23,7 +23,8 @@ """ import json -from typing import Union, Dict, Any, Tuple +from collections import defaultdict +from typing import Union, Dict, Any, Tuple, List import pandas as pd from hbutils.design import SingletonMark @@ -81,7 +82,7 @@ def _open_onnx_model(model_name: str): @ts_lru_cache() -def _open_tags(model_name: str) -> pd.DataFrame: +def _open_tags(model_name: str) -> Tuple[pd.DataFrame, Dict[str, List[str]]]: """ Load the tag metadata from Hugging Face Hub with caching. @@ -95,11 +96,18 @@ def _open_tags(model_name: str) -> pd.DataFrame: :return: DataFrame containing tag information with columns like 'name', 'category' :rtype: pd.DataFrame """ - return pd.read_csv(hf_hub_download( + df_tags = pd.read_csv(hf_hub_download( repo_id=_get_repo_id(model_name), repo_type='model', filename='selected_tags.csv', )) + d_ips = {} + if 'ips' in df_tags: + df_tags['ips'] = df_tags['ips'].map(json.loads) + for name, ips in zip(df_tags['name'], df_tags['ips']): + if ips: + d_ips[name] = ips + return df_tags, d_ips @ts_lru_cache() @@ -250,7 +258,7 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', >>> import numpy as np >>> normalized_embedding = embedding / np.linalg.norm(embedding) """ - df_tags = _open_tags(model_name=model_name) + df_tags, d_ips = _open_tags(model_name=model_name) values = _raw_predict(image, model_name=model_name) prediction = values['prediction'] tags = {} @@ -288,4 +296,13 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', tags.update(cate_tags) values['tag'] = tags + ip_mapping, ip_counts = {}, defaultdict(lambda: 0) + if 'ips' in df_tags.columns: + for tag, _ in tags.items(): + if tag in d_ips[tag]: + ip_mapping[tag] = d_ips[tag] + for ip_name in d_ips[tag]: + ip_counts[ip_name] += 1 + values['ips_mapping'] = ip_mapping + values['ips'] = [x for x, _ in sorted(ip_counts.items(), key=lambda x: (-x[1], x[0]))] return vreplace(fmt, values) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index f16ed4ef6f0..bfa5018dcaf 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -60,7 +60,7 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): exts.append([]) else: exts.append([]) - df_tags['ips'] = exts + df_tags['ips'] = list(map(json.dumps, exts)) logging.info(f'Tags:\n{df_tags}') df_tags.to_csv(os.path.join(upload_dir, 'selected_tags.csv'), index=False) From 561af44773473efd583a5afe1b911fb05095e157 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 11:16:24 +0800 Subject: [PATCH 19/32] dev(narugo): adjust some values detail --- imgutils/tagging/pixai.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py index 8fada7289ab..5b26077c94e 100644 --- a/imgutils/tagging/pixai.py +++ b/imgutils/tagging/pixai.py @@ -296,13 +296,14 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', tags.update(cate_tags) values['tag'] = tags - ip_mapping, ip_counts = {}, defaultdict(lambda: 0) + ips_mapping, ips_counts = {}, defaultdict(lambda: 0) if 'ips' in df_tags.columns: for tag, _ in tags.items(): - if tag in d_ips[tag]: - ip_mapping[tag] = d_ips[tag] + if tag in d_ips: + ips_mapping[tag] = d_ips[tag] for ip_name in d_ips[tag]: - ip_counts[ip_name] += 1 - values['ips_mapping'] = ip_mapping - values['ips'] = [x for x, _ in sorted(ip_counts.items(), key=lambda x: (-x[1], x[0]))] + ips_counts[ip_name] += 1 + values['ips_mapping'] = ips_mapping + values['ips_count'] = dict(ips_counts) + values['ips'] = [x for x, _ in sorted(ips_counts.items(), key=lambda x: (-x[1], x[0]))] return vreplace(fmt, values) From 47fc3859a7e636198897df94a38f0e342defecbc Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 11:20:33 +0800 Subject: [PATCH 20/32] dev(narugo): fix the docs --- imgutils/tagging/pixai.py | 50 +++++++++++++++++++++++++++++---------- 1 file changed, 38 insertions(+), 12 deletions(-) diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py index 5b26077c94e..adb74748e26 100644 --- a/imgutils/tagging/pixai.py +++ b/imgutils/tagging/pixai.py @@ -9,6 +9,11 @@ versions of these models for efficient inference, available at `deepghs `_ repositories. + In addition to standard tagging, the models can identify anime character IP (Intellectual Property) + associations. For example, if a character like "misaka_mikoto" is detected, the system can map + this to the "toaru_kagaku_no_railgun" (A Certain Scientific Railgun) IP. All IP names follow + Danbooru-style tag conventions for consistency with existing anime tagging systems. + Example:: >>> from imgutils.tagging.pixai import get_pixai_tags >>> # Get tags with default thresholds @@ -20,6 +25,10 @@ >>> # Get all tags in a single dictionary >>> all_tags = get_pixai_tags('path/to/image.jpg', fmt='tag') >>> print("All tags:", all_tags) + + >>> # Get IP information for detected characters + >>> ips = get_pixai_tags('path/to/image.jpg', fmt='ips') + >>> print("Detected IPs:", ips) """ import json @@ -86,15 +95,16 @@ def _open_tags(model_name: str) -> Tuple[pd.DataFrame, Dict[str, List[str]]]: """ Load the tag metadata from Hugging Face Hub with caching. - This function downloads and loads the CSV file containing tag names and categories - for the specified model. The DataFrame contains columns for tag names, categories, - and other metadata. + This function downloads and loads the CSV file containing tag names, categories, + and IP (Intellectual Property) associations for the specified model. The DataFrame + contains columns for tag names, categories, and other metadata including character + IP mappings when available. :param model_name: Name of the model :type model_name: str - :return: DataFrame containing tag information with columns like 'name', 'category' - :rtype: pd.DataFrame + :return: Tuple containing (DataFrame with tag information, dictionary mapping character tags to their IPs) + :rtype: Tuple[pd.DataFrame, Dict[str, List[str]]] """ df_tags = pd.read_csv(hf_hub_download( repo_id=_get_repo_id(model_name), @@ -212,7 +222,7 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', :type image: ImageTyping :param model_name: Name or path of the PixAI tagger model to use :type model_name: str - :param thresholds: Confidence threshold values. Can be a single float applied to all + :param thresholds: Confidence threshold values. Can be a single float applied to all categories, or a dictionary mapping category IDs/names to specific thresholds :type thresholds: Union[float, Dict[Any, float]], optional :param fmt: Output format specification. If FMT_UNSET, returns all available categories. @@ -226,37 +236,53 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', .. note:: The fmt parameter can include the following keys: - - Category names (e.g., 'general', 'character'): dictionaries containing category-specific + - Category names (e.g., 'general', 'character'): dictionaries containing category-specific tags and their confidence scores - ``tag``: a dictionary containing all tags across categories and their confidences - - ``embedding``: a 1-dimensional embedding vector of the image, recommended for similarity + - ``embedding``: a 1-dimensional embedding vector of the image, recommended for similarity search after L2 normalization - ``logits``: raw 1-dimensional logits output from the model - ``prediction``: 1-dimensional prediction probabilities from the model + - ``ips_mapping``: a dictionary mapping detected character tags to their associated + IP (Intellectual Property) names in Danbooru-style format + - ``ips_count``: a dictionary containing IP names and their occurrence counts based + on detected characters + - ``ips``: a list of IP names sorted by occurrence count (descending) and name + (ascending), representing the most likely anime/game series in the image Default category thresholds are used if not specified. These vary by model and category but typically range from 0.35 to 0.5. Example:: >>> from imgutils.tagging.pixai import get_pixai_tags - >>> + >>> >>> # Get tags with default format (all categories) >>> general_tags, character_tags = get_pixai_tags('anime_image.jpg', model_name='v0.9') >>> print("General tags:", general_tags) >>> print("Character tags:", character_tags) - >>> + >>> >>> # Get all tags in a single dictionary >>> all_tags = get_pixai_tags('image.jpg', fmt='tag') >>> print("All tags:", all_tags) - >>> + >>> >>> # Use custom thresholds >>> result = get_pixai_tags('image.jpg', thresholds={'general': 0.3, 'character': 0.5}) - >>> + >>> >>> # Get embedding for similarity search >>> embedding = get_pixai_tags('image.jpg', fmt='embedding') >>> # Normalize for cosine similarity >>> import numpy as np >>> normalized_embedding = embedding / np.linalg.norm(embedding) + >>> + >>> # Get IP information for character identification + >>> ips_mapping = get_pixai_tags('image.jpg', fmt='ips_mapping') + >>> print("Character to IP mapping:", ips_mapping) + >>> # Example output: {'misaka_mikoto': ['toaru_kagaku_no_railgun'], 'hu_tao_(genshin_impact)': ['genshin_impact']} + >>> + >>> # Get most likely anime/game series + >>> top_ips = get_pixai_tags('image.jpg', fmt='ips') + >>> print("Most likely series:", top_ips) + >>> # Example output: ['genshin_impact', 'toaru_kagaku_no_railgun'] """ df_tags, d_ips = _open_tags(model_name=model_name) values = _raw_predict(image, model_name=model_name) From 744b7a720af92c2ed1db485460fbe684e0c2e5dc Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 11:38:02 +0800 Subject: [PATCH 21/32] dev(narugo): add unittest for it --- test/tagging/test_pixai.py | 173 +++++++++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 test/tagging/test_pixai.py diff --git a/test/tagging/test_pixai.py b/test/tagging/test_pixai.py new file mode 100644 index 00000000000..afb32d20391 --- /dev/null +++ b/test/tagging/test_pixai.py @@ -0,0 +1,173 @@ +import numpy as np +import pytest + +from imgutils.tagging import get_pixai_tags +from imgutils.tagging.pixai import _open_tags, _open_preprocess, _open_onnx_model, _open_default_category_thresholds +from test.testings import get_testfile + + +@pytest.fixture(autouse=True, scope='module') +def _release_model_after_run(): + try: + yield + finally: + _open_tags.cache_clear() + _open_preprocess.cache_clear() + _open_onnx_model.cache_clear() + _open_default_category_thresholds.cache_clear() + + +@pytest.mark.unittest +class TestTaggingPixAI: + def test_get_pixai_tags(self): + tags, chars = get_pixai_tags(get_testfile('6124220.jpg')) + + assert tags['cat_girl'] >= 0.8 + assert not chars + assert isinstance(tags['cat_girl'], float) + + tags, chars = get_pixai_tags(get_testfile('6125785.jpg')) + assert tags['1girl'] >= 0.90 + assert chars['hu_tao_(genshin_impact)'] >= 0.95 + assert isinstance(tags['1girl'], float) + assert isinstance(chars['hu_tao_(genshin_impact)'], float) + + def test_pixai_tags_sample(self): + tags, chars = get_pixai_tags(get_testfile('6125785.jpg')) + + assert tags == pytest.approx({ + 'symbol-shaped_pupils': 0.9938011169433594, 'ghost': 0.9909012317657471, + 'flower-shaped_pupils': 0.9810856580734253, 'porkpie_hat': 0.9693692922592163, + 'black_nails': 0.9692082405090332, 'ghost_pose': 0.9579154253005981, 'ring': 0.9509532451629639, + '1girl': 0.9435760974884033, 'hat_flower': 0.9315004348754883, 'hat': 0.9308193922042847, + 'flower': 0.9203469753265381, 'plum_blossoms': 0.8814681768417358, 'red_shirt': 0.8734415769577026, + 'jewelry': 0.8436187505722046, 'claw_pose': 0.8382970094680786, 'chinese_clothes': 0.8213527202606201, + 'long_sleeves': 0.8101956844329834, 'long_hair': 0.7956619262695312, 'twintails': 0.7865841388702393, + 'looking_at_viewer': 0.7795377373695374, 'hat_tassel': 0.7775378227233887, + 'multiple_rings': 0.7483937740325928, 'signature': 0.7478364109992981, 'smile': 0.741124153137207, + 'open_mouth': 0.7302751541137695, 'red_flower': 0.705192506313324, 'solo': 0.6761863827705383, + 'hat_ornament': 0.6489882469177246, 'red_eyes': 0.6290297508239746, 'black_hat': 0.6113986372947693, + 'brown_coat': 0.6067467927932739, 'nail_polish': 0.5937596559524536, 'coat': 0.5664999485015869, + 'upper_body': 0.5661808848381042, 'brown_hair': 0.5366549491882324, 'thumb_ring': 0.5287110209465027, + 'star_(symbol)': 0.5028221011161804, 'shirt': 0.4754156470298767, 'star-shaped_pupils': 0.45203927159309387, + 'hair_between_eyes': 0.4512987434864044, 'orange_eyes': 0.4298587739467621, ':d': 0.4219313859939575, + 'tangzhuang': 0.3894977569580078, ':3': 0.3866003453731537, 'blush': 0.3761531412601471, + 'v-shaped_eyebrows': 0.34583818912506104, 'sidelocks': 0.3450319468975067, + 'brown_shirt': 0.31902527809143066, 'blurry': 0.3167526423931122 + }, abs=2e-2) + assert chars == pytest.approx({ + 'hu_tao_(genshin_impact)': 0.9999773502349854, 'boo_tao_(genshin_impact)': 0.9996482133865356 + }, abs=2e-2) + + def test_pixai_rgba(self): + tags, chars = get_pixai_tags(get_testfile('nian.png')) + assert tags == pytest.approx({ + 'red_tube_top': 0.9998065233230591, 'tube_top': 0.9994766712188721, 'red_bandeau': 0.9992805123329163, + 'white_shorts': 0.9976104497909546, 'bandeau': 0.9968563914299011, + 'transparent_background': 0.9929141402244568, 'horns': 0.9903750419616699, 'full_body': 0.9860043525695801, + 'midriff': 0.9834859371185303, 'shorts': 0.9828792810440063, '1girl': 0.982780933380127, + 'strapless': 0.9817174673080444, 'navel': 0.9786539673805237, 'white_footwear': 0.9764545559883118, + 'white_hair': 0.9713669419288635, 'solo': 0.9699287414550781, 'tail': 0.9654039740562439, + 'stomach': 0.9624348282814026, 'standing': 0.9548828601837158, 'pointy_ears': 0.9536501169204712, + 'open_coat': 0.9320983290672302, 'looking_at_viewer': 0.9311813116073608, + 'open_clothes': 0.9311613440513611, 'long_hair': 0.9210299253463745, 'hand_on_own_hip': 0.9110379219055176, + 'holding': 0.9101808071136475, 'short_shorts': 0.9012227058410645, 'tongue': 0.8910117745399475, + 'tongue_out': 0.8871811628341675, 'dragon_horns': 0.8710488080978394, 'weapon': 0.8667179346084595, + 'coat': 0.8642098307609558, 'belt': 0.8602883219718933, 'tachi-e': 0.8591709136962891, + 'crop_top': 0.856552004814148, 'white_coat': 0.8258588314056396, 'dragon_tail': 0.822641909122467, + 'purple_eyes': 0.7998863458633423, 'streaked_hair': 0.796852707862854, 'half_updo': 0.790149450302124, + 'thighs': 0.7805066108703613, 'bare_legs': 0.7744182348251343, ':d': 0.7707932591438293, + 'wide_sleeves': 0.767501711845398, 'smile': 0.765789270401001, 'multicolored_hair': 0.7620936632156372, + 'bead_bracelet': 0.7511173486709595, 'red_hair': 0.7497332692146301, 'drop_shadow': 0.7467425465583801, + 'long_sleeves': 0.7251672744750977, 'red_skin': 0.7151302099227905, 'sidelocks': 0.6788227558135986, + 'socks': 0.6740190982818604, ':p': 0.6698598861694336, 'sword': 0.6676492691040039, + 'jewelry': 0.6632057428359985, 'hand_up': 0.637663722038269, 'boots': 0.6372023224830627, + 'black_socks': 0.626619815826416, 'pixel_art': 0.6166024208068848, 'black_belt': 0.6042207479476929, + 'holding_weapon': 0.6014618873596191, 'breasts': 0.6006667017936707, 'ankle_boots': 0.5968820452690125, + 'white_jacket': 0.5880962014198303, 'jacket': 0.5767049789428711, 'tassel_earrings': 0.5749104022979736, + 'open_mouth': 0.5724694132804871, 'beads': 0.5716468095779419, 'braid': 0.5678756237030029, + 'open_jacket': 0.5464628338813782, 'bare_shoulders': 0.5238108038902283, 'shoes': 0.506024181842804, + 'bracelet': 0.4898264706134796, 'pouch': 0.48912930488586426, 'colored_skin': 0.46792706847190857, + 'originium_arts_(arknights)': 0.4517717957496643, 'dragon_girl': 0.4428599774837494, + 'thigh_strap': 0.4376051723957062, 'earrings': 0.4229695200920105, 'collarbone': 0.41589391231536865, + 'red_eyes': 0.4031139016151428, 'body_markings': 0.4004215598106384, 'medium_breasts': 0.384432315826416, + 'holding_sword': 0.3541863262653351, 'white_pants': 0.34769535064697266, + 'club_(weapon)': 0.33692681789398193, 'eyeshadow': 0.3357437252998352, 'ponytail': 0.3333626389503479, + 'tassel': 0.33294352889060974, 'multiple_weapons': 0.3174567222595215, 'pink_eyes': 0.31579816341400146, + 'flame-tipped_tail': 0.3054245710372925 + }, abs=2e-2) + assert chars == pytest.approx({'nian_(arknights)': 0.9999774694442749}, abs=2e-2) + + def test_pixai_tags_sample_ips(self): + tags, chars, ips, ips_count, ips_mapping = get_pixai_tags( + get_testfile('6125785.jpg'), + fmt=('general', 'character', 'ips', 'ips_count', 'ips_mapping') + ) + + assert tags == pytest.approx({ + 'symbol-shaped_pupils': 0.9938011169433594, 'ghost': 0.9909012317657471, + 'flower-shaped_pupils': 0.9810856580734253, 'porkpie_hat': 0.9693692922592163, + 'black_nails': 0.9692082405090332, 'ghost_pose': 0.9579154253005981, 'ring': 0.9509532451629639, + '1girl': 0.9435760974884033, 'hat_flower': 0.9315004348754883, 'hat': 0.9308193922042847, + 'flower': 0.9203469753265381, 'plum_blossoms': 0.8814681768417358, 'red_shirt': 0.8734415769577026, + 'jewelry': 0.8436187505722046, 'claw_pose': 0.8382970094680786, 'chinese_clothes': 0.8213527202606201, + 'long_sleeves': 0.8101956844329834, 'long_hair': 0.7956619262695312, 'twintails': 0.7865841388702393, + 'looking_at_viewer': 0.7795377373695374, 'hat_tassel': 0.7775378227233887, + 'multiple_rings': 0.7483937740325928, 'signature': 0.7478364109992981, 'smile': 0.741124153137207, + 'open_mouth': 0.7302751541137695, 'red_flower': 0.705192506313324, 'solo': 0.6761863827705383, + 'hat_ornament': 0.6489882469177246, 'red_eyes': 0.6290297508239746, 'black_hat': 0.6113986372947693, + 'brown_coat': 0.6067467927932739, 'nail_polish': 0.5937596559524536, 'coat': 0.5664999485015869, + 'upper_body': 0.5661808848381042, 'brown_hair': 0.5366549491882324, 'thumb_ring': 0.5287110209465027, + 'star_(symbol)': 0.5028221011161804, 'shirt': 0.4754156470298767, 'star-shaped_pupils': 0.45203927159309387, + 'hair_between_eyes': 0.4512987434864044, 'orange_eyes': 0.4298587739467621, ':d': 0.4219313859939575, + 'tangzhuang': 0.3894977569580078, ':3': 0.3866003453731537, 'blush': 0.3761531412601471, + 'v-shaped_eyebrows': 0.34583818912506104, 'sidelocks': 0.3450319468975067, + 'brown_shirt': 0.31902527809143066, 'blurry': 0.3167526423931122 + }, abs=2e-2) + assert chars == pytest.approx({ + 'hu_tao_(genshin_impact)': 0.9999773502349854, 'boo_tao_(genshin_impact)': 0.9996482133865356 + }, abs=2e-2) + assert ips == ['genshin_impact'] + assert ips_count == {'genshin_impact': 2} + assert ips_mapping == { + 'hu_tao_(genshin_impact)': ['genshin_impact'], + 'boo_tao_(genshin_impact)': ['genshin_impact'], + } + + def test_pixai_tags_sample_custom(self): + tags, chars, embedding, logits, prediction = get_pixai_tags( + get_testfile('6125785.jpg'), + fmt=('general', 'character', 'embedding', 'logits', 'prediction'), + model_name='deepghs/pixai-tagger-v0.9-onnx', + ) + + assert tags == pytest.approx({ + 'symbol-shaped_pupils': 0.9938011169433594, 'ghost': 0.9909012317657471, + 'flower-shaped_pupils': 0.9810856580734253, 'porkpie_hat': 0.9693692922592163, + 'black_nails': 0.9692082405090332, 'ghost_pose': 0.9579154253005981, 'ring': 0.9509532451629639, + '1girl': 0.9435760974884033, 'hat_flower': 0.9315004348754883, 'hat': 0.9308193922042847, + 'flower': 0.9203469753265381, 'plum_blossoms': 0.8814681768417358, 'red_shirt': 0.8734415769577026, + 'jewelry': 0.8436187505722046, 'claw_pose': 0.8382970094680786, 'chinese_clothes': 0.8213527202606201, + 'long_sleeves': 0.8101956844329834, 'long_hair': 0.7956619262695312, 'twintails': 0.7865841388702393, + 'looking_at_viewer': 0.7795377373695374, 'hat_tassel': 0.7775378227233887, + 'multiple_rings': 0.7483937740325928, 'signature': 0.7478364109992981, 'smile': 0.741124153137207, + 'open_mouth': 0.7302751541137695, 'red_flower': 0.705192506313324, 'solo': 0.6761863827705383, + 'hat_ornament': 0.6489882469177246, 'red_eyes': 0.6290297508239746, 'black_hat': 0.6113986372947693, + 'brown_coat': 0.6067467927932739, 'nail_polish': 0.5937596559524536, 'coat': 0.5664999485015869, + 'upper_body': 0.5661808848381042, 'brown_hair': 0.5366549491882324, 'thumb_ring': 0.5287110209465027, + 'star_(symbol)': 0.5028221011161804, 'shirt': 0.4754156470298767, 'star-shaped_pupils': 0.45203927159309387, + 'hair_between_eyes': 0.4512987434864044, 'orange_eyes': 0.4298587739467621, ':d': 0.4219313859939575, + 'tangzhuang': 0.3894977569580078, ':3': 0.3866003453731537, 'blush': 0.3761531412601471, + 'v-shaped_eyebrows': 0.34583818912506104, 'sidelocks': 0.3450319468975067, + 'brown_shirt': 0.31902527809143066, 'blurry': 0.3167526423931122 + }, abs=2e-2) + assert chars == pytest.approx({ + 'hu_tao_(genshin_impact)': 0.9999773502349854, 'boo_tao_(genshin_impact)': 0.9996482133865356 + }, abs=2e-2) + + assert isinstance(embedding, np.ndarray) + assert embedding.shape == (1024,) + assert isinstance(logits, np.ndarray) + assert logits.shape == (13461,) + assert isinstance(prediction, np.ndarray) + assert prediction.shape == (13461,) From e7d67cd5b39051a0fde600182c2bc86e7ce9a05a Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 11:42:04 +0800 Subject: [PATCH 22/32] dev(narugo): add benchmark --- .../api_doc/tagging/tagging_benchmark.plot.py | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/docs/source/api_doc/tagging/tagging_benchmark.plot.py b/docs/source/api_doc/tagging/tagging_benchmark.plot.py index e110be999bc..d01d287836d 100644 --- a/docs/source/api_doc/tagging/tagging_benchmark.plot.py +++ b/docs/source/api_doc/tagging/tagging_benchmark.plot.py @@ -2,7 +2,7 @@ from benchmark import BaseBenchmark, create_plot_cli from imgutils.tagging import get_deepdanbooru_tags, get_wd14_tags, get_mldanbooru_tags, get_deepgelbooru_tags, \ - get_camie_tags + get_camie_tags, get_pixai_tags class DeepdanbooruBenchmark(BaseBenchmark): @@ -95,6 +95,32 @@ def run(self): _ = get_camie_tags(image_file, model_name=self.model_name) +class PixAITaggerBenchmark(BaseBenchmark): + def __init__(self, model_name: str): + BaseBenchmark.__init__(self) + self.model_name = model_name + + def load(self): + from imgutils.tagging.pixai import _open_tags, _open_preprocess, _open_onnx_model, \ + _open_default_category_thresholds + _ = _open_tags(self.model_name) + _ = _open_preprocess(self.model_name) + _ = _open_onnx_model(self.model_name) + _ = _open_default_category_thresholds(self.model_name) + + def unload(self): + from imgutils.tagging.pixai import _open_tags, _open_preprocess, _open_onnx_model, \ + _open_default_category_thresholds + _open_tags.cache_clear() + _open_preprocess.cache_clear() + _open_onnx_model.cache_clear() + _open_default_category_thresholds.cache_clear() + + def run(self): + image_file = random.choice(self.all_images) + _ = get_pixai_tags(image_file, model_name=self.model_name) + + if __name__ == '__main__': create_plot_cli( [ @@ -113,6 +139,7 @@ def run(self): ('mldanbooru', MLDanbooruBenchmark()), ('camie-initial', CamieBenchmark('initial')), ('camie-refined', CamieBenchmark('refined')), + ('pixai-tagger-v0.9', PixAITaggerBenchmark('v0.9')), ], title='Benchmark for Tagging Models', run_times=10, From e1ff174d1a624a89217a0bcdc2fe479acc227184 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 12:25:00 +0800 Subject: [PATCH 23/32] dev(narugo): add cache deletion on CI environment --- .../tagging/tagging_benchmark.plot.py.svg | 3315 ----------------- test/conftest.py | 9 + 2 files changed, 9 insertions(+), 3315 deletions(-) delete mode 100644 docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg diff --git a/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg b/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg deleted file mode 100644 index d86cc8a7681..00000000000 --- a/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg +++ /dev/null @@ -1,3315 +0,0 @@ - - - - - - - - 2025-03-29T17:54:03.129689 - image/svg+xml - - - Matplotlib v3.7.5, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/test/conftest.py b/test/conftest.py index d4f615a88d0..12188f64285 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,5 +1,8 @@ +import os + import pytest from hbutils.testing import TextAligner +from hfutils.cache import delete_cache try: import torch @@ -15,3 +18,9 @@ def _try_import_torch(): @pytest.fixture() def text_aligner(): return TextAligner().multiple_lines() + + +@pytest.fixture(autouse=True, scope='module') +def clean_hf_cache(): + if os.environ.get('CI'): + delete_cache() From e7b8cd15eb73e48f0f5d325725afb61b62fc710b Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 14:33:32 +0800 Subject: [PATCH 24/32] dev(narugo): try fix unittest --- docs/source/_libs/benchmark.py | 4 ++ .../api_doc/tagging/tagging_benchmark.plot.py | 37 +++++++++++++------ test/conftest.py | 7 +++- 3 files changed, 34 insertions(+), 14 deletions(-) diff --git a/docs/source/_libs/benchmark.py b/docs/source/_libs/benchmark.py index 7f7bf3c4430..d2d02a54f92 100644 --- a/docs/source/_libs/benchmark.py +++ b/docs/source/_libs/benchmark.py @@ -34,6 +34,9 @@ def load(self): def unload(self): raise NotImplementedError + def after_unload(self): + pass + def run(self): raise NotImplementedError @@ -60,6 +63,7 @@ def _record(name): self.unload() _record('') + self.after_unload() mems = np.array([mem for _, mem, _ in logs]) mems -= mems[0] diff --git a/docs/source/api_doc/tagging/tagging_benchmark.plot.py b/docs/source/api_doc/tagging/tagging_benchmark.plot.py index d01d287836d..c51328409fc 100644 --- a/docs/source/api_doc/tagging/tagging_benchmark.plot.py +++ b/docs/source/api_doc/tagging/tagging_benchmark.plot.py @@ -1,25 +1,34 @@ import random +from hfutils.cache import delete_cache + from benchmark import BaseBenchmark, create_plot_cli from imgutils.tagging import get_deepdanbooru_tags, get_wd14_tags, get_mldanbooru_tags, get_deepgelbooru_tags, \ get_camie_tags, get_pixai_tags -class DeepdanbooruBenchmark(BaseBenchmark): +class CleanModelStorageBenchmark(BaseBenchmark): + def after_unload(self): + delete_cache() + + +class DeepdanbooruBenchmark(CleanModelStorageBenchmark): def load(self): - from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model + from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model, _get_deepdanbooru_labels _ = _get_deepdanbooru_model() + _ = _get_deepdanbooru_labels def unload(self): - from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model + from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model, _get_deepdanbooru_labels _get_deepdanbooru_model.cache_clear() + _get_deepdanbooru_labels.cache_clear() def run(self): image_file = random.choice(self.all_images) _ = get_deepdanbooru_tags(image_file) -class DeepgelbooruBenchmark(BaseBenchmark): +class DeepgelbooruBenchmark(CleanModelStorageBenchmark): def load(self): from imgutils.tagging.deepgelbooru import _open_tags, _open_model, _open_preprocessor _ = _open_tags() @@ -37,39 +46,43 @@ def run(self): _ = get_deepgelbooru_tags(image_file) -class Wd14Benchmark(BaseBenchmark): +class Wd14Benchmark(CleanModelStorageBenchmark): def __init__(self, model): BaseBenchmark.__init__(self) self.model = model def load(self): - from imgutils.tagging.wd14 import _get_wd14_model + from imgutils.tagging.wd14 import _get_wd14_model, _get_wd14_labels _ = _get_wd14_model(self.model) + _ = _get_wd14_labels(self.model) def unload(self): - from imgutils.tagging.wd14 import _get_wd14_model + from imgutils.tagging.wd14 import _get_wd14_model, _get_wd14_labels _get_wd14_model.cache_clear() + _get_wd14_labels.cache_clear() def run(self): image_file = random.choice(self.all_images) _ = get_wd14_tags(image_file, model_name=self.model) -class MLDanbooruBenchmark(BaseBenchmark): +class MLDanbooruBenchmark(CleanModelStorageBenchmark): def load(self): - from imgutils.tagging.mldanbooru import _open_mldanbooru_model + from imgutils.tagging.mldanbooru import _open_mldanbooru_model, _get_mldanbooru_labels _ = _open_mldanbooru_model() + _ = _get_mldanbooru_labels() def unload(self): - from imgutils.tagging.mldanbooru import _open_mldanbooru_model + from imgutils.tagging.mldanbooru import _open_mldanbooru_model, _get_mldanbooru_labels _open_mldanbooru_model.cache_clear() + _get_mldanbooru_labels.cache_clear() def run(self): image_file = random.choice(self.all_images) _ = get_mldanbooru_tags(image_file) -class CamieBenchmark(BaseBenchmark): +class CamieBenchmark(CleanModelStorageBenchmark): def __init__(self, model_name): BaseBenchmark.__init__(self) self.model_name = model_name @@ -95,7 +108,7 @@ def run(self): _ = get_camie_tags(image_file, model_name=self.model_name) -class PixAITaggerBenchmark(BaseBenchmark): +class PixAITaggerBenchmark(CleanModelStorageBenchmark): def __init__(self, model_name: str): BaseBenchmark.__init__(self) self.model_name = model_name diff --git a/test/conftest.py b/test/conftest.py index 12188f64285..1b93873f6e3 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -22,5 +22,8 @@ def text_aligner(): @pytest.fixture(autouse=True, scope='module') def clean_hf_cache(): - if os.environ.get('CI'): - delete_cache() + try: + yield + finally: + if os.environ.get('CI'): + delete_cache() From 138ada24df046b0afd31e5628611bfd2bc47eab1 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 14:50:46 +0800 Subject: [PATCH 25/32] dev(narugo): try fix unittest, ci skip --- .github/workflows/doc.yml | 2 ++ docs/source/api_doc/tagging/tagging_benchmark.plot.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/doc.yml b/.github/workflows/doc.yml index 3417574ba98..2df2b2a4dab 100644 --- a/.github/workflows/doc.yml +++ b/.github/workflows/doc.yml @@ -67,6 +67,7 @@ jobs: ENV_PROD: 'true' PLANTUML_HOST: http://localhost:18080 HF_TOKEN: ${{ secrets.HF_TOKEN }} + CI: 'true' run: | plantumlcli -c make docs @@ -145,6 +146,7 @@ jobs: ENV_PROD: 'true' PLANTUML_HOST: http://localhost:18080 HF_TOKEN: ${{ secrets.HF_TOKEN }} + CI: 'true' run: | git fetch --all --tags git branch -av diff --git a/docs/source/api_doc/tagging/tagging_benchmark.plot.py b/docs/source/api_doc/tagging/tagging_benchmark.plot.py index c51328409fc..b81ebf47b90 100644 --- a/docs/source/api_doc/tagging/tagging_benchmark.plot.py +++ b/docs/source/api_doc/tagging/tagging_benchmark.plot.py @@ -1,3 +1,4 @@ +import os import random from hfutils.cache import delete_cache @@ -9,7 +10,8 @@ class CleanModelStorageBenchmark(BaseBenchmark): def after_unload(self): - delete_cache() + if os.environ.get('CI'): + delete_cache() class DeepdanbooruBenchmark(CleanModelStorageBenchmark): From 16f0cd9eb774d8bbc14839df9613ebe3b0af49ce Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 15:08:28 +0800 Subject: [PATCH 26/32] dev(narugo): try fix unittest, ci skip --- zoo/pixai_tagger/export.py | 42 +++++++++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index bfa5018dcaf..caf77c293f0 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -1,5 +1,8 @@ import json import os.path +import re +from pprint import pformat +from textwrap import indent import numpy as np import onnx @@ -7,7 +10,7 @@ import pandas as pd import torch from ditk import logging -from hbutils.string import titleize +from hbutils.string import titleize, underscore from hbutils.system import TemporaryDirectory from hbutils.testing import vpip from hfutils.operate import get_hf_client, upload_directory_as_directory @@ -18,13 +21,14 @@ from imgutils.data import load_image from imgutils.preprocess import parse_torchvision_transforms +from imgutils.tagging import get_pixai_tags from zoo.pixai_tagger.tags import load_tags from zoo.utils import onnx_optimize, get_testfile, torch_model_profile_via_calflops from .min_script import EndpointHandler from .onnx import get_model -def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): +def sync(src_repo: str, dst_repo: str, no_optimize: bool = False, show_current_results: bool = False): hf_client = get_hf_client() if not hf_client.repo_exists(repo_id=dst_repo, repo_type='model'): hf_client.create_repo(repo_id=dst_repo, repo_type='model', private=True) @@ -212,6 +216,34 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): print(f'```', file=f) print(f'', file=f) + matching = re.fullmatch(r'^deepghs/pixai-tagger-(?P[\s\S]+)-onnx$', dst_repo) + model_name = matching.group('version') if matching else dst_repo + print(f'```python', file=f) + print(f'from imgutils.tagging import get_pixai_tags', file=f) + print(f'', file=f) + + cate_names = tuple((cate_name for _, cate_name in sorted(d_category_names.items()))) + fmt_names = tuple([*cate_names, 'ips', 'ips_mapping']) + var_names = tuple(map(underscore, fmt_names)) + print(f'{", ".join(var_names)} = get_pixai_tags(', file=f) + print(f' {sample_input_url!r},', file=f) + print(f' model_name={model_name!r},', file=f) + print(f' fmt={fmt_names!r},', file=f) + print(f')', file=f) + print(f'', file=f) + + var_values = get_pixai_tags(sample_input_file, model_name=model_name, fmt=fmt_names) + for varname, fmtname, varvalue in zip(var_names, fmt_names, var_values): + print(f'print({varname})', file=f) + print(indent( + pformat(varvalue, sort_dicts=False), + prefix='# ', + ), file=f) + + print(f'', file=f) + print(f'```', file=f) + print(f'', file=f) + upload_directory_as_directory( repo_id=dst_repo, repo_type='model', @@ -231,7 +263,11 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False): if __name__ == '__main__': logging.try_init_root(logging.INFO) + # matching = re.fullmatch(r'^deepghs/pixai-tagger-(?P[\s\S]+)-onnx$', 'deepghs/pixai-tagger-v0.9-onnx') + # print(matching) + # print(matching.group('version')) sync( src_repo='pixai-labs/pixai-tagger-v0.9', - dst_repo='deepghs/pixai-tagger-v0.9-onnx' + dst_repo='deepghs/pixai-tagger-v0.9-onnx', + show_current_results=True, ) From 4b895d6d9811046ddc354aab27b3a4cfba49477f Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 15:18:04 +0800 Subject: [PATCH 27/32] dev(narugo): try fix unittest, ci skip --- zoo/pixai_tagger/export.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index caf77c293f0..983058eb831 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -189,9 +189,22 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False, show_current_r f'{len(df_tags[df_tags["category"] == category])}', file=f) print(f'', file=f) - print(f'## How to Use', file=f) + print(f'## Thresholds', file=f) + print(f'', file=f) + ths = [] + for item in df_th.to_dict('records'): + ths.append({ + 'Category': item['category'], + 'Name': item['name'], + 'Count': len(df_tags[df_tags['category'] == item['category']]), + 'Threshold': item['threshold'], + }) + df_th_shown = pd.DataFrame(ths) + print(df_th_shown.to_markdown(index=False), file=f) print(f'', file=f) + print(f'## How to Use', file=f) + print(f'', file=f) imgutils_version = str(vpip('dghs-imgutils')._actual_version) sample_input = dummy_image if min(sample_input.width, sample_input.height) > 640: From 41c1a096fc3b34aacf0b957b9c6d1519f2dea0b5 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 16:30:07 +0800 Subject: [PATCH 28/32] dev(narugo): try ignore those 2 unittests --- test/generic/test_classify.py | 3 +++ test/generic/test_yolo.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/test/generic/test_classify.py b/test/generic/test_classify.py index ed789b940f6..06ccad9d0cc 100644 --- a/test/generic/test_classify.py +++ b/test/generic/test_classify.py @@ -1,8 +1,10 @@ +from unittest import skipUnless from unittest.mock import patch import numpy as np import pytest from PIL import Image +from hbutils.testing import vpip from huggingface_hub.utils import reset_sessions from imgutils.generic import classify_predict_score @@ -232,6 +234,7 @@ def test_classify_predict_fmt_complex(self): assert np.linalg.norm(results['embedding']) == pytest.approx(np.linalg.norm(expected_embedding)) @patch("huggingface_hub.constants.HF_HUB_OFFLINE", True) + @skipUnless(vpip('huggingface_hub') < '0.34', 'Has problem on huggingface 0.34+') def test_classify_predict_score_top5_offline_mode(self, clean_session): image = Image.open(get_testfile('png_640.png')) scores = classify_predict_score( diff --git a/test/generic/test_yolo.py b/test/generic/test_yolo.py index d472af2fb16..ba9bf341635 100644 --- a/test/generic/test_yolo.py +++ b/test/generic/test_yolo.py @@ -1,6 +1,8 @@ +from unittest import skipUnless from unittest.mock import patch import pytest +from hbutils.testing import vpip from huggingface_hub import configure_http_backend from huggingface_hub.utils import reset_sessions @@ -52,6 +54,7 @@ def test_detect_faces_none(self): ) == [] @patch("huggingface_hub.constants.HF_HUB_OFFLINE", True) + @skipUnless(vpip('huggingface_hub') < '0.34', 'Has problem on huggingface 0.34+') def test_detect_faces_with_offline_mode(self, clean_session): configure_http_backend() detection = yolo_predict( From 57b797a4061822f4794801c8c0c9bac48e8ed70b Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 10:30:51 +0000 Subject: [PATCH 29/32] dev(narugo): auto sync Tue, 09 Sep 2025 10:30:51 +0000 --- .../tagging/tagging_benchmark.plot.py.svg | 3445 +++++++++++++++++ 1 file changed, 3445 insertions(+) create mode 100644 docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg diff --git a/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg b/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg new file mode 100644 index 00000000000..071703ca40f --- /dev/null +++ b/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg @@ -0,0 +1,3445 @@ + + + + + + + + 2025-09-09T10:30:22.529066 + image/svg+xml + + + Matplotlib v3.7.5, https://matplotlib.org/ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + From fff687197114bda60c6d4c4f4eba1b2eabffa3ce Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 22:56:53 +0800 Subject: [PATCH 30/32] dev(narugo): try fix this unittest --- Makefile | 2 +- imgutils/tagging/pixai.py | 21 ++++++++++++++++----- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/Makefile b/Makefile index 49fe1f7d5a9..044dbff10c9 100644 --- a/Makefile +++ b/Makefile @@ -41,7 +41,7 @@ unittest: --cov="${RANGE_SRC_DIR}" \ $(if ${MIN_COVERAGE},--cov-fail-under=${MIN_COVERAGE},) \ $(if ${WORKERS},-n ${WORKERS},) \ - --reruns 8 --reruns-delay 2 --only-rerun '(OSError|Timeout|HTTPError.*429|HTTPError.*502|HTTPError.*504|check your connection|429 error)' + --reruns 8 --reruns-delay 2 --only-rerun '(OSError|Timeout|HTTPError.*429|HTTPError.*502|HTTPError.*504|check your connection|429 error|CAS service error)' docs: $(MAKE) -C "${DOC_DIR}" build diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py index adb74748e26..b2a620fff4a 100644 --- a/imgutils/tagging/pixai.py +++ b/imgutils/tagging/pixai.py @@ -220,7 +220,7 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', :param image: The input image to analyze (file path, PIL Image, numpy array, etc.) :type image: ImageTyping - :param model_name: Name or path of the PixAI tagger model to use + :param model_name: Name or repository ID of the PixAI tagger model to use :type model_name: str :param thresholds: Confidence threshold values. Can be a single float applied to all categories, or a dictionary mapping category IDs/names to specific thresholds @@ -253,6 +253,17 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', Default category thresholds are used if not specified. These vary by model and category but typically range from 0.35 to 0.5. + You can extract embedding of the given image with the following code + + >>> from imgutils.tagging import get_pixai_tags + >>> + >>> embedding = get_pixai_tags('skadi.jpg', fmt='embedding') + >>> embedding.shape + (1024, ) + + This embedding is valuable for constructing indices that enable rapid querying of images based on + visual features within large-scale datasets. + Example:: >>> from imgutils.tagging.pixai import get_pixai_tags >>> @@ -322,14 +333,14 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', tags.update(cate_tags) values['tag'] = tags - ips_mapping, ips_counts = {}, defaultdict(lambda: 0) if 'ips' in df_tags.columns: + ips_mapping, ips_counts = {}, defaultdict(lambda: 0) for tag, _ in tags.items(): if tag in d_ips: ips_mapping[tag] = d_ips[tag] for ip_name in d_ips[tag]: ips_counts[ip_name] += 1 - values['ips_mapping'] = ips_mapping - values['ips_count'] = dict(ips_counts) - values['ips'] = [x for x, _ in sorted(ips_counts.items(), key=lambda x: (-x[1], x[0]))] + values['ips_mapping'] = ips_mapping + values['ips_count'] = dict(ips_counts) + values['ips'] = [x for x, _ in sorted(ips_counts.items(), key=lambda x: (-x[1], x[0]))] return vreplace(fmt, values) From 760ddcec91c27bd20bb87b73821df6b0587fc955 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 23:39:16 +0800 Subject: [PATCH 31/32] dev(narugo): add actual usage docs --- imgutils/tagging/pixai.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py index b2a620fff4a..2063bc3eb3d 100644 --- a/imgutils/tagging/pixai.py +++ b/imgutils/tagging/pixai.py @@ -294,6 +294,23 @@ def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', >>> top_ips = get_pixai_tags('image.jpg', fmt='ips') >>> print("Most likely series:", top_ips) >>> # Example output: ['genshin_impact', 'toaru_kagaku_no_railgun'] + + Here are some images for example + + .. image:: tagging_demo.plot.py.svg + :align: center + + >>> general, character = get_pixai_tags('skadi.jpg') + >>> general + {'patreon_username': 0.9988852739334106, 'baseball_bat': 0.9977256059646606, 'holding_baseball_bat': 0.9858889579772949, 'navel': 0.9830228090286255, 'crop_top': 0.9666315317153931, 'sportswear': 0.9664723873138428, '1girl': 0.9572311639785767, 'long_hair': 0.9550737738609314, 'outdoors': 0.9501817226409912, 'solo': 0.9466996788978577, 'day': 0.9394471049308777, 'breasts': 0.938787579536438, 'web_address': 0.9387772679328918, 'stomach': 0.935083270072937, 'red_eyes': 0.9326196908950806, 'shorts': 0.9305683374404907, 'motion_blur': 0.9278550148010254, 'playing_sports': 0.9263769388198853, 'blue_sky': 0.9213213920593262, 'midriff': 0.9191423654556274, 'large_breasts': 0.9174768924713135, 'artist_name': 0.9089528322219849, 'sky': 0.9054281711578369, 'baseball': 0.904181957244873, 'gloves': 0.9033604860305786, 'thighs': 0.893738865852356, 'black_shorts': 0.8926981687545776, 'volleyball': 0.8198539614677429, 'very_long_hair': 0.7967187166213989, 'short_shorts': 0.7873305082321167, 'black_gloves': 0.7765249013900757, 'white_hair': 0.770541787147522, 'baseball_mitt': 0.7684446573257446, 'thigh_gap': 0.73811936378479, 'sweat': 0.7263807654380798, 'cowboy_shot': 0.7235408425331116, 'short_sleeves': 0.7062878012657166, 'parted_lips': 0.7025120258331299, 'patreon_logo': 0.6970672607421875, 'cloud': 0.6967148780822754, 'looking_at_viewer': 0.6898926496505737, 'holding': 0.6879030466079712, 'swinging': 0.6736525893211365, 'ass_visible_through_thighs': 0.6636734008789062, 'elbow_pads': 0.6630151867866516, 'shirt': 0.6250661611557007, 'hair_between_eyes': 0.6075361967086792, 'standing': 0.5285079479217529, 'black_shirt': 0.5173394680023193, 'linea_alba': 0.513701319694519, 'baseball_uniform': 0.48175835609436035, 'crop_top_overhang': 0.4682246744632721, 'ball': 0.43616628646850586, 'blurry': 0.4201475977897644, 'baseball_stadium': 0.41493287682533264, 'grey_hair': 0.39384859800338745, 'watermark': 0.3919041156768799, 'black_sports_bra': 0.3877854645252228, 'fanbox_username': 0.3790855407714844, 'narrow_waist': 0.36392998695373535} + >>> character + {'skadi_(arknights)': 0.8926791548728943} + >>> + >>> general, character = get_pixai_tags('hutao.jpg') + >>> general + {'bag': 0.9833353757858276, 'backpack': 0.9766197204589844, 'flower-shaped_pupils': 0.962916910648346, 'tongue_out': 0.960152804851532, 'tongue': 0.9526823163032532, 'ghost': 0.9514724016189575, 'plaid_skirt': 0.9499615430831909, '1girl': 0.9378864765167236, 'skirt': 0.9353114366531372, 'bag_charm': 0.9314961433410645, 'symbol-shaped_pupils': 0.9252510070800781, 'charm_(object)': 0.9249529242515564, 'twintails': 0.9239017367362976, 'flower': 0.9175764322280884, 'outdoors': 0.9175151586532593, 'hair_ornament': 0.9161680936813354, 'plaid_clothes': 0.9144806861877441, 'long_hair': 0.8768749833106995, 'pleated_skirt': 0.8597153425216675, 'school_uniform': 0.8573414087295532, 'looking_at_viewer': 0.8392735719680786, ':p': 0.8193913698196411, 'hair_between_eyes': 0.8070638179779053, 'hair_flower': 0.8054562211036682, 'nail_polish': 0.8011300563812256, 'building': 0.7961824536323547, 'jacket': 0.7647742629051208, 'brown_hair': 0.7541910409927368, 'solo': 0.7539198398590088, 'long_sleeves': 0.7471930980682373, 'ahoge': 0.7171378135681152, 'hair_ribbon': 0.6994943618774414, 'red_eyes': 0.6819245219230652, 'bowtie': 0.6639955043792725, 'sidelocks': 0.6275356411933899, 'bush': 0.6164096593856812, 'gate': 0.612525224685669, 'smile': 0.6077383160591125, 'shirt': 0.6042965650558472, 'contemporary': 0.5968752503395081, 'brick_floor': 0.5933602452278137, 'cardigan': 0.5810519456863403, 'gradient_hair': 0.5570307970046997, 'diagonal-striped_bowtie': 0.5565160512924194, 'alternate_costume': 0.5535630583763123, 'school_bag': 0.5535626411437988, 'black_hair': 0.5434530973434448, 'ribbon': 0.5332301259040833, 'hairclip': 0.523446261882782, 'day': 0.5164296627044678, 'street': 0.49987316131591797, 'bow': 0.4941294193267822, 'plum_blossoms': 0.4940766990184784, 'collared_shirt': 0.49013280868530273, 'standing': 0.4820355772972107, 'blue_cardigan': 0.47836002707481384, 'cowboy_shot': 0.4782080054283142, 'pocket': 0.477585107088089, 'pavement': 0.4712265729904175, 'multicolored_hair': 0.4610708951950073, 'blue_jacket': 0.45271334052085876, 'blush': 0.45005902647972107, 'sleeves_past_wrists': 0.440649151802063, 'black_nails': 0.4402858316898346, 'black_bag': 0.4206739366054535, 'miniskirt': 0.4187837243080139, 'red_bow': 0.414681613445282, 'very_long_hair': 0.4129619002342224, 'diagonal-striped_clothes': 0.4112803339958191, 'blazer': 0.40750616788864136, 'striped_bowtie': 0.40123170614242554, 'sunlight': 0.4008329212665558, 'grey_skirt': 0.3930213749408722, 'road': 0.3819067180156708, 'black_ribbon': 0.3776353895664215, 'thighs': 0.3722286820411682, 'hug': 0.37215015292167664, 'brick_wall': 0.3717171549797058, 'white_shirt': 0.3694952428340912, 'open_clothes': 0.36442798376083374, 'open_jacket': 0.3525886535644531, ':d': 0.3343055844306946, 'multicolored_nails': 0.32190075516700745, 'red_bowtie': 0.3157669007778168, 'star-shaped_pupils': 0.309164822101593, 'open_mouth': 0.30890953540802, 'beads': 0.3084579110145569, 'stone_stairs': 0.30559882521629333, 'randoseru': 0.30517613887786865} + >>> character + {'hu_tao_(genshin_impact)': 0.9997367858886719, 'boo_tao_(genshin_impact)': 0.999537467956543} """ df_tags, d_ips = _open_tags(model_name=model_name) values = _raw_predict(image, model_name=model_name) From e9af3868af9a4d00e982a670a8c9200a81a75da2 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Tue, 9 Sep 2025 23:42:44 +0800 Subject: [PATCH 32/32] dev(narugo): add one more line for export function --- zoo/pixai_tagger/export.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py index 983058eb831..e3be359a608 100644 --- a/zoo/pixai_tagger/export.py +++ b/zoo/pixai_tagger/export.py @@ -229,6 +229,11 @@ def sync(src_repo: str, dst_repo: str, no_optimize: bool = False, show_current_r print(f'```', file=f) print(f'', file=f) + print(f'You can use function `get_pixai_tags` to tag your image. ' + f'For more details of this function, see [documentation of imgutils.tagging.pixai]' + f'(https://dghs-imgutils.deepghs.org/main/api_doc/tagging/pixai.html).', file=f) + print(f'', file=f) + matching = re.fullmatch(r'^deepghs/pixai-tagger-(?P[\s\S]+)-onnx$', dst_repo) model_name = matching.group('version') if matching else dst_repo print(f'```python', file=f)