From aaaa713f01e6f9ad6e09f6be69398b2d466efd61 Mon Sep 17 00:00:00 2001 From: corkborg <62742115+corkborg@users.noreply.github.com> Date: Tue, 25 Mar 2025 20:03:38 +0900 Subject: [PATCH] Implement camie-tagger. --- .../interrogator/camietaggerinterrogator.py | 154 +++++++++--------- tagger/interrogators.py | 7 + 2 files changed, 85 insertions(+), 76 deletions(-) diff --git a/tagger/interrogator/camietaggerinterrogator.py b/tagger/interrogator/camietaggerinterrogator.py index e2ce668..6fe4b0c 100644 --- a/tagger/interrogator/camietaggerinterrogator.py +++ b/tagger/interrogator/camietaggerinterrogator.py @@ -1,30 +1,22 @@ -import os +""" +Interrogator for the Camie-Tagger model. +reference: +https://huggingface.co/Camais03/camie-tagger/blob/main/onnx_inference.py +""" import sys -import pandas as pd -import numpy as np - -from typing import Iterable, Tuple, List, Dict -from PIL import Image - -from pathlib import Path -from huggingface_hub import hf_hub_download -import re import json -import torchvision.transforms as transforms +from typing import cast -from numpy import asarray, float32, expand_dims, exp +import numpy as np +import torch +from PIL import Image +from huggingface_hub import hf_hub_download +from onnxruntime import InferenceSession -tag_escape_pattern = re.compile(r'([\\()])') +from tagger.interrogator.interrogator import AbsInterrogator -import tagger.dbimutils as dbimutils - -class CamieTaggerInterrogator(Interrogator): - """ - Interrogator for the Camie-Tagger model. - reference: - https://huggingface.co/Camais03/camie-tagger/blob/main/onnx_inference.py - """ +class CamieTaggerInterrogator(AbsInterrogator): repo_id: str model_path: str tags_path: str @@ -34,7 +26,7 @@ class CamieTaggerInterrogator(Interrogator): name: str, repo_id: str, model_path: str, - tags_path='classes.json', + tags_path='metadata.json', ) -> None: super().__init__(name) self.model_path = model_path @@ -43,7 +35,7 @@ class CamieTaggerInterrogator(Interrogator): self.tags = None self.model = None - def download(self) -> Tuple[str, str]: + def download(self) -> tuple[str, str]: print(f"Loading {self.name} model file from {self.repo_id}", file=sys.stderr) model_path = hf_hub_download( @@ -59,83 +51,93 @@ class CamieTaggerInterrogator(Interrogator): def load(self) -> None: model_path, tags_path = self.download() - from onnxruntime import InferenceSession self.model = InferenceSession(model_path, providers=self.providers) print(f'Loaded {self.name} model from {model_path}', file=sys.stderr) with open(tags_path, 'r', encoding='utf-8') as filen: - self.tags = json.load(filen) + self.metadata = json.load(filen) def interrogate( self, image: Image.Image - ) -> Tuple[ - Dict[str, float], # rating confidents - Dict[str, float] # tag confidents + ) -> tuple[ + dict[str, float], # rating confidents + dict[str, float] # tag confidents ]: # init model if self.model is None: self.load() + if self.model is None: + raise Exception("Model not loading.") - image = dbimutils.fill_transparent(image) - image = dbimutils.resize(image, 448) # TODO CUSTOMIZE - - x = asarray(image, dtype=float32) / 255 - # HWC -> 1CHW - x = x.transpose((2, 0, 1)) - x = expand_dims(x, 0) + img_tensor = preprocess_image(image) + img_numpy = img_tensor.unsqueeze(0).numpy() input_ = self.model.get_inputs()[0] - output = self.model.get_outputs()[0] + # evaluate model - y, = self.model.run([output.name], {input_.name: x}) + outputs = self.model.run(None, {input_.name: img_numpy}) - # Softmax - y = 1 / (1 + exp(-y)) + # Process outputs + initial_probs: np.ndarray = 1.0 / (1.0 + np.exp(-outputs[0])) # Apply sigmoid + refined_probs: np.ndarray = 1.0 / (1.0 + np.exp(-outputs[1])) if len(outputs) > 1 else initial_probs - tags = {tag: float(conf) for tag, conf in zip(self.tags, y.flatten())} - return {}, tags + # Get top tags + indices = np.atleast_1d(refined_probs[0]).nonzero()[0] -def preprocess_image(image_path, image_size=512): + # Group by category + tags_by_category = {} + for idx in indices: + idx_str = str(idx) + tag_name = self.metadata['idx_to_tag'].get(idx_str, f"unknown-{idx}") + category = self.metadata['tag_to_category'].get(tag_name, "general") + + if category not in tags_by_category: + tags_by_category[category] = [] + + prob = float(refined_probs[0, idx]) + tags_by_category[category].append((tag_name, prob)) + + # 'year', 'rating', 'general', 'character', 'copyright', 'artist', 'meta' + return dict(tags_by_category['rating']), dict(tags_by_category['general']) + +def preprocess_image(img: Image.Image, image_size=512) -> torch.Tensor: """Process an image for inference""" - if not os.path.exists(image_path): - raise ValueError(f"Image not found at path: {image_path}") + + # Convert RGBA or Palette images to RGB + if img.mode in ('RGBA', 'P'): + img = img.convert('RGB') + + # Get original dimensions + width, height = img.size + aspect_ratio = width / height + + # Calculate new dimensions to maintain aspect ratio + if aspect_ratio > 1: + new_width = image_size + new_height = int(new_width / aspect_ratio) + else: + new_height = image_size + new_width = int(new_height * aspect_ratio) + + # Resize with LANCZOS filter + img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) + + # Create new image with padding + new_image = Image.new('RGB', (image_size, image_size), (0, 0, 0)) + paste_x = (image_size - new_width) // 2 + paste_y = (image_size - new_height) // 2 + new_image.paste(img, (paste_x, paste_y)) + + import torchvision.transforms as transforms # Initialize transform transform = transforms.Compose([ transforms.ToTensor(), ]) - try: - with Image.open(image_path) as img: - # Convert RGBA or Palette images to RGB - if img.mode in ('RGBA', 'P'): - img = img.convert('RGB') - - # Get original dimensions - width, height = img.size - aspect_ratio = width / height - - # Calculate new dimensions to maintain aspect ratio - if aspect_ratio > 1: - new_width = image_size - new_height = int(new_width / aspect_ratio) - else: - new_height = image_size - new_width = int(new_height * aspect_ratio) - - # Resize with LANCZOS filter - img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) - - # Create new image with padding - new_image = Image.new('RGB', (image_size, image_size), (0, 0, 0)) - paste_x = (image_size - new_width) // 2 - paste_y = (image_size - new_height) // 2 - new_image.paste(img, (paste_x, paste_y)) - - # Apply transforms - img_tensor = transform(new_image) - return img_tensor - except Exception as e: - raise Exception(f"Error processing {image_path}: {str(e)}") + # Apply transforms + img_tensor= transform(new_image) + img_tensor = cast(torch.Tensor, img_tensor) + return img_tensor diff --git a/tagger/interrogators.py b/tagger/interrogators.py index 763ec03..6882f5a 100644 --- a/tagger/interrogators.py +++ b/tagger/interrogators.py @@ -1,6 +1,7 @@ from tagger.interrogator.interrogator import AbsInterrogator from tagger.interrogator.waifudiffusioninterrogator import WaifuDiffusionInterrogator from tagger.interrogator.mldanbooruinterrogator import MLDanbooruInterrogator +from tagger.interrogator.camietaggerinterrogator import CamieTaggerInterrogator interrogators: dict[str, AbsInterrogator] = { 'wd14-vit.v1': WaifuDiffusionInterrogator( @@ -73,4 +74,10 @@ interrogators: dict[str, AbsInterrogator] = { repo_id='deepghs/ml-danbooru-onnx', model_path='TResnet-D-FLq_ema_6-30000.onnx' ), + 'camie-tagger': CamieTaggerInterrogator( + 'Camie Tagger', + repo_id='Camais03/camie-tagger', + model_path='model_initial.onnx', + tags_path='model_initial_metadata.json' + ), }