diff --git a/run.py b/run.py index df42da4..a0a1818 100644 --- a/run.py +++ b/run.py @@ -1,5 +1,7 @@ +import sys + from typing import Generator, Iterable -from tagger.interrogator import Interrogator +from tagger.interrogator.interrogator import AbsInterrogator from PIL import Image from pathlib import Path import argparse @@ -79,7 +81,7 @@ def image_interrogate(image_path: Path, tag_escape: bool, exclude_tags: Iterable im = Image.open(image_path) result = interrogator.interrogate(im) - return Interrogator.postprocess_tags( + return AbsInterrogator.postprocess_tags( result[1], threshold=args.threshold, escape_tag=tag_escape, @@ -116,7 +118,7 @@ if args.dir: if args.file: tags = image_interrogate(Path(args.file), not args.rawtag, parse_exclude_tags()) - print() + print(file=sys.stderr) tags_str = ', '.join(tags.keys()) print(tags_str) diff --git a/tagger/interrogator.py b/tagger/interrogator.py deleted file mode 100644 index 4f1d42a..0000000 --- a/tagger/interrogator.py +++ /dev/null @@ -1,256 +0,0 @@ -import os -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 - -from numpy import asarray, float32, expand_dims, exp - -tag_escape_pattern = re.compile(r'([\\()])') - -import tagger.dbimutils as dbimutils - - -class Interrogator: - @staticmethod - def postprocess_tags( - tags: Dict[str, float], - threshold=0.35, - additional_tags: List[str] = [], - exclude_tags: Iterable[str] = [], - sort_by_alphabetical_order=False, - add_confident_as_weight=False, - replace_underscore=False, - replace_underscore_excludes: List[str] = [], - escape_tag=False - ) -> Dict[str, float]: - for t in additional_tags: - tags[t] = 1.0 - - # those lines are totally not "pythonic" but looks better to me - tags = { - t: c - - # sort by tag name or confident - for t, c in sorted( - tags.items(), - key=lambda i: i[0 if sort_by_alphabetical_order else 1], - reverse=not sort_by_alphabetical_order - ) - - # filter tags - if ( - c >= threshold - and t not in exclude_tags - ) - } - - new_tags = [] - for tag in list(tags): - new_tag = tag - - if replace_underscore and tag not in replace_underscore_excludes: - new_tag = new_tag.replace('_', ' ') - - if escape_tag: - new_tag = tag_escape_pattern.sub(r'\\\1', new_tag) - - if add_confident_as_weight: - new_tag = f'({new_tag}:{tags[tag]})' - - new_tags.append((new_tag, tags[tag])) - tags = dict(new_tags) - - return tags - - def __init__(self, name: str) -> None: - self.name = name - self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] - - def load(self): - raise NotImplementedError() - - def unload(self) -> bool: - unloaded = False - - if hasattr(self, 'model') and self.model is not None: - del self.model - unloaded = True - print(f'Unloaded {self.name}') - - if hasattr(self, 'tags'): - del self.tags - - return unloaded - - def use_cpu(self) -> None: - self.providers = ['CPUExecutionProvider'] - - def interrogate( - self, - image: Image - ) -> Tuple[ - Dict[str, float], # rating confidents - Dict[str, float] # tag confidents - ]: - raise NotImplementedError() - -class WaifuDiffusionInterrogator(Interrogator): - def __init__( - self, - name: str, - model_path='model.onnx', - tags_path='selected_tags.csv', - **kwargs - ) -> None: - super().__init__(name) - self.model_path = model_path - self.tags_path = tags_path - self.kwargs = kwargs - - def download(self) -> Tuple[os.PathLike, os.PathLike]: - print(f"Loading {self.name} model file from {self.kwargs['repo_id']}") - - model_path = Path(hf_hub_download( - **self.kwargs, filename=self.model_path)) - tags_path = Path(hf_hub_download( - **self.kwargs, filename=self.tags_path)) - return model_path, tags_path - - def load(self) -> None: - model_path, tags_path = self.download() - - from onnxruntime import InferenceSession - self.model = InferenceSession(str(model_path), providers=self.providers) - - print(f'Loaded {self.name} model from {model_path}') - - self.tags = pd.read_csv(tags_path) - - def interrogate( - self, - image: Image - ) -> Tuple[ - Dict[str, float], # rating confidents - Dict[str, float] # tag confidents - ]: - # init model - if not hasattr(self, 'model') or self.model is None: - self.load() - - # code for converting the image and running the model is taken from the link below - # thanks, SmilingWolf! - # https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags/blob/main/app.py - - # convert an image to fit the model - _, height, _, _ = self.model.get_inputs()[0].shape - - # alpha to white - image = image.convert('RGBA') - new_image = Image.new('RGBA', image.size, 'WHITE') - new_image.paste(image, mask=image) - image = new_image.convert('RGB') - image = np.asarray(image) - - # PIL RGB to OpenCV BGR - image = image[:, :, ::-1] - - image = dbimutils.make_square(image, height) - image = dbimutils.smart_resize(image, height) - image = image.astype(np.float32) - image = np.expand_dims(image, 0) - - # evaluate model - input_name = self.model.get_inputs()[0].name - label_name = self.model.get_outputs()[0].name - confidents = self.model.run([label_name], {input_name: image})[0] - - tags = self.tags[:][['name']] - tags['confidents'] = confidents[0] - - # first 4 items are for rating (general, sensitive, questionable, explicit) - ratings = dict(tags[:4].values) - - # rest are regular tags - tags = dict(tags[4:].values) - - return ratings, tags - -class MLDanbooruInterrogator(Interrogator): - """ Interrogator for the MLDanbooru model. """ - def __init__( - self, - name: str, - repo_id: str, - model_path: str, - tags_path='classes.json', - ) -> None: - super().__init__(name) - self.model_path = model_path - self.tags_path = tags_path - self.repo_id = repo_id - self.tags = None - self.model = None - - def download(self) -> Tuple[str, str]: - print(f"Loading {self.name} model file from {self.repo_id}") - - model_path = hf_hub_download( - repo_id=self.repo_id, - filename=self.model_path - ) - tags_path = hf_hub_download( - repo_id=self.repo_id, - filename=self.tags_path, - ) - return model_path, tags_path - - 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}') - - with open(tags_path, 'r', encoding='utf-8') as filen: - self.tags = json.load(filen) - - def interrogate( - self, - image: Image - ) -> Tuple[ - Dict[str, float], # rating confidents - Dict[str, float] # tag confidents - ]: - # init model - if self.model is None: - self.load() - - 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) - - input_ = self.model.get_inputs()[0] - output = self.model.get_outputs()[0] - # evaluate model - y, = self.model.run([output.name], {input_.name: x}) - - # Softmax - y = 1 / (1 + exp(-y)) - - tags = {tag: float(conf) for tag, conf in zip(self.tags, y.flatten())} - return {}, tags - - def large_batch_interrogate(self, images: List, dry_run=False) -> str: - raise NotImplementedError() diff --git a/tagger/interrogator/__init__.py b/tagger/interrogator/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tagger/interrogator/camietaggerinterrogator.py b/tagger/interrogator/camietaggerinterrogator.py new file mode 100644 index 0000000..e2ce668 --- /dev/null +++ b/tagger/interrogator/camietaggerinterrogator.py @@ -0,0 +1,141 @@ +import os +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 numpy import asarray, float32, expand_dims, exp + +tag_escape_pattern = re.compile(r'([\\()])') + +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 + """ + repo_id: str + model_path: str + tags_path: str + + def __init__( + self, + name: str, + repo_id: str, + model_path: str, + tags_path='classes.json', + ) -> None: + super().__init__(name) + self.model_path = model_path + self.tags_path = tags_path + self.repo_id = repo_id + self.tags = None + self.model = None + + 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( + repo_id=self.repo_id, + filename=self.model_path + ) + tags_path = hf_hub_download( + repo_id=self.repo_id, + filename=self.tags_path, + ) + return model_path, tags_path + + 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) + + def interrogate( + self, + image: Image.Image + ) -> Tuple[ + Dict[str, float], # rating confidents + Dict[str, float] # tag confidents + ]: + # init model + if self.model is None: + self.load() + + 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) + + input_ = self.model.get_inputs()[0] + output = self.model.get_outputs()[0] + # evaluate model + y, = self.model.run([output.name], {input_.name: x}) + + # Softmax + y = 1 / (1 + exp(-y)) + + tags = {tag: float(conf) for tag, conf in zip(self.tags, y.flatten())} + return {}, tags + +def preprocess_image(image_path, image_size=512): + """Process an image for inference""" + if not os.path.exists(image_path): + raise ValueError(f"Image not found at path: {image_path}") + + # 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)}") diff --git a/tagger/interrogator/interrogator.py b/tagger/interrogator/interrogator.py new file mode 100644 index 0000000..3e968b9 --- /dev/null +++ b/tagger/interrogator/interrogator.py @@ -0,0 +1,97 @@ +import re +import sys + +import pandas as pd + +from typing import Iterable, Tuple, List, Dict +from PIL import Image + +from onnxruntime import InferenceSession + +tag_escape_pattern = re.compile(r'([\\()])') + +class AbsInterrogator: + model: InferenceSession | None + tags: pd.DataFrame | None + @staticmethod + def postprocess_tags( + tags: Dict[str, float], + threshold=0.35, + additional_tags: List[str] = [], + exclude_tags: Iterable[str] = [], + sort_by_alphabetical_order=False, + add_confident_as_weight=False, + replace_underscore=False, + replace_underscore_excludes: List[str] = [], + escape_tag=False + ) -> Dict[str, float]: + for t in additional_tags: + tags[t] = 1.0 + + # those lines are totally not "pythonic" but looks better to me + tags = { + t: c + + # sort by tag name or confident + for t, c in sorted( + tags.items(), + key=lambda i: i[0 if sort_by_alphabetical_order else 1], + reverse=not sort_by_alphabetical_order + ) + + # filter tags + if ( + c >= threshold + and t not in exclude_tags + ) + } + + new_tags = [] + for tag in list(tags): + new_tag = tag + + if replace_underscore and tag not in replace_underscore_excludes: + new_tag = new_tag.replace('_', ' ') + + if escape_tag: + new_tag = tag_escape_pattern.sub(r'\\\1', new_tag) + + if add_confident_as_weight: + new_tag = f'({new_tag}:{tags[tag]})' + + new_tags.append((new_tag, tags[tag])) + tags = dict(new_tags) + + return tags + + def __init__(self, name: str) -> None: + self.name = name + self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] + + def load(self): + raise NotImplementedError() + + def unload(self) -> bool: + unloaded = False + + if hasattr(self, 'model') and self.model is not None: + del self.model + unloaded = True + print(f'Unloaded {self.name}', file=sys.stderr) + + if hasattr(self, 'tags'): + del self.tags + + return unloaded + + def use_cpu(self) -> None: + self.providers = ['CPUExecutionProvider'] + + def interrogate( + self, + image: Image.Image + ) -> Tuple[ + Dict[str, float], # rating confidents + Dict[str, float] # tag confidents + ]: + raise NotImplementedError() diff --git a/tagger/interrogator/mldanbooruinterrogator.py b/tagger/interrogator/mldanbooruinterrogator.py new file mode 100644 index 0000000..49f4ba0 --- /dev/null +++ b/tagger/interrogator/mldanbooruinterrogator.py @@ -0,0 +1,87 @@ +import sys +import json + +from PIL import Image +from numpy import asarray, float32, expand_dims, exp + +from huggingface_hub import hf_hub_download + +from tagger.interrogator.interrogator import AbsInterrogator +import tagger.dbimutils as dbimutils + +class MLDanbooruInterrogator(AbsInterrogator): + """ Interrogator for the MLDanbooru model. """ + def __init__( + self, + name: str, + repo_id: str, + model_path: str, + tags_path='classes.json', + ) -> None: + super().__init__(name) + self.model_path = model_path + self.tags_path = tags_path + self.repo_id = repo_id + self.tags = None + self.model = None + + 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( + repo_id=self.repo_id, + filename=self.model_path + ) + tags_path = hf_hub_download( + repo_id=self.repo_id, + filename=self.tags_path, + ) + return model_path, tags_path + + 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) + + def interrogate( + self, + image: Image.Image + ) -> tuple[ + dict[str, float], # rating confidents + dict[str, float] # tag confidents + ]: + # init model + if self.model is None: + self.load() + + 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) + + input_ = self.model.get_inputs()[0] + output = self.model.get_outputs()[0] + # evaluate model + y, = self.model.run([output.name], {input_.name: x}) + + # Softmax + y = 1 / (1 + exp(-y)) + + if self.tags is None: + raise Exception("Tags not loading.") + + tags = {tag: float(conf) for tag, conf in zip(self.tags, y.flatten())} + return {}, tags + + def large_batch_interrogate(self, images: list, dry_run=False) -> str: + raise NotImplementedError() + diff --git a/tagger/interrogator/waifudiffusioninterrogator.py b/tagger/interrogator/waifudiffusioninterrogator.py new file mode 100644 index 0000000..ebe7a42 --- /dev/null +++ b/tagger/interrogator/waifudiffusioninterrogator.py @@ -0,0 +1,103 @@ +import os +import sys +import pandas as pd +import numpy as np + +from typing import Tuple, Dict +from PIL import Image + +from pathlib import Path +from huggingface_hub import hf_hub_download +import re + +tag_escape_pattern = re.compile(r'([\\()])') + +from tagger.interrogator.interrogator import AbsInterrogator +import tagger.dbimutils as dbimutils + +class WaifuDiffusionInterrogator(AbsInterrogator): + def __init__( + self, + name: str, + model_path='model.onnx', + tags_path='selected_tags.csv', + **kwargs + ) -> None: + super().__init__(name) + self.model_path = model_path + self.tags_path = tags_path + self.kwargs = kwargs + + def download(self) -> Tuple[os.PathLike, os.PathLike]: + print(f"Loading {self.name} model file from {self.kwargs['repo_id']}", file=sys.stderr) + + model_path = Path(hf_hub_download( + **self.kwargs, filename=self.model_path)) + tags_path = Path(hf_hub_download( + **self.kwargs, filename=self.tags_path)) + return model_path, tags_path + + def load(self) -> None: + model_path, tags_path = self.download() + + from onnxruntime import InferenceSession + self.model = InferenceSession(str(model_path), providers=self.providers) + + print(f'Loaded {self.name} model from {model_path}', file=sys.stderr) + + self.tags = pd.read_csv(tags_path) + + def interrogate( + self, + input_image: Image.Image + ) -> tuple[ + dict[str, float], # rating confidents + dict[str, float] # tag confidents + ]: + # init model + if not hasattr(self, 'model') or self.model is None: + self.load() + + if self.model is None: + raise Exception("Model not loading.") + + # code for converting the image and running the model is taken from the link below + # thanks, SmilingWolf! + # https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags/blob/main/app.py + + # convert an image to fit the model + _, height, _, _ = self.model.get_inputs()[0].shape + + # alpha to white + image = input_image.convert('RGBA') + new_image = Image.new('RGBA', image.size, 'WHITE') + new_image.paste(image, mask=image) + image = new_image.convert('RGB') + image = np.asarray(image) + + # PIL RGB to OpenCV BGR + image = image[:, :, ::-1] + + image = dbimutils.make_square(image, height) + image = dbimutils.smart_resize(image, height) + image = image.astype(np.float32) + image = np.expand_dims(image, 0) + + # evaluate model + input_name = self.model.get_inputs()[0].name + label_name = self.model.get_outputs()[0].name + confidents = self.model.run([label_name], {input_name: image})[0] + + if self.tags is None: + raise Exception("Tags not loading.") + + tags = self.tags[:][['name']] + tags['confidents'] = confidents[0] + + # first 4 items are for rating (general, sensitive, questionable, explicit) + ratings = dict(tags[:4].values) + + # rest are regular tags + tags = dict(tags[4:].values) + + return ratings, tags diff --git a/tagger/interrogators.py b/tagger/interrogators.py index 01c18ce..763ec03 100644 --- a/tagger/interrogators.py +++ b/tagger/interrogators.py @@ -1,8 +1,8 @@ -from typing import List, Dict +from tagger.interrogator.interrogator import AbsInterrogator +from tagger.interrogator.waifudiffusioninterrogator import WaifuDiffusionInterrogator +from tagger.interrogator.mldanbooruinterrogator import MLDanbooruInterrogator -from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator, MLDanbooruInterrogator - -interrogators: Dict[str, Interrogator] = { +interrogators: dict[str, AbsInterrogator] = { 'wd14-vit.v1': WaifuDiffusionInterrogator( 'WD14 ViT v1', repo_id='SmilingWolf/wd-v1-4-vit-tagger'