From ba6c3b640a605f291559fc740d90452f6955a100 Mon Sep 17 00:00:00 2001 From: corkborg <62742115+corkborg@users.noreply.github.com> Date: Thu, 28 Dec 2023 02:28:15 +0900 Subject: [PATCH] Support multi models. --- README.md | 26 ++++++++++--- run.py | 18 +++++---- tagger/dbimutils.py | 20 ++++++++++ tagger/interrogator.py | 83 +++++++++++++++++++++++++++++++++++++++-- tagger/interrogators.py | 46 +++++++++++++++++++++++ 5 files changed, 176 insertions(+), 17 deletions(-) create mode 100644 tagger/interrogators.py diff --git a/README.md b/README.md index 198dcf1..11142f5 100644 --- a/README.md +++ b/README.md @@ -5,16 +5,16 @@ forked from [https://github.com/picobyte/stable-diffusion-webui-wd14-tagger]() ## install ``` -venv/bin/python -m venv venv +python -m venv venv source venv/bin/activate pip install -r requirements.txt - ``` ## usage ``` -usage: run.py [-h] [--dir DIR | --file FILE] [--threthold THRETHOLD] [--ext EXT] +usage: run.py [-h] (--dir DIR | --file FILE) [--threthold THRETHOLD] [--ext EXT] + [--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}] options: -h, --help show this help message and exit @@ -23,16 +23,32 @@ options: --threthold THRETHOLD 予測値の足切り確率（デフォルトは0.35） --ext EXT dirの場合にキャプションファイルにつける拡張子 + --model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000} + 予測に使用するモデル名 ``` single file ``` -python3 run.py --file image.jpg +python run.py --file image.jpg ``` batch execution ``` -python3 run.py --dir dir/dir +python run.py --dir dir/dir +``` + +## Support Models + +``` +python run.py --file image.jpg --model wd14-vit.v1 +python run.py --file image.jpg --model wd14-vit.v2 +python run.py --file image.jpg --model wd14-convnext.v1 +python run.py --file image.jpg --model wd14-convnext.v2 +python run.py --file image.jpg --model wd14-convnextv2.v1 +python run.py --file image.jpg --model wd14-swinv2-v1 +python run.py --file image.jpg --model wd-v1-4-moat-tagger.v2 +python run.py --file image.jpg --model mld-caformer.dec-5-97527 +python run.py --file image.jpg --model mld-tresnetd.6-30000 ``` \ No newline at end of file diff --git a/run.py b/run.py index d61c8cb..ba000a9 100644 --- a/run.py +++ b/run.py @@ -6,9 +6,11 @@ from PIL import Image from pathlib import Path import argparse +from tagger.interrogators import interrogators + parser = argparse.ArgumentParser() -group = parser.add_mutually_exclusive_group() +group = parser.add_mutually_exclusive_group(required=True) group.add_argument('--dir', help='ディレクトリ内の画像すべてに予測を行う') group.add_argument('--file', help='ファイルに対して予測を行う') @@ -21,16 +23,16 @@ parser.add_argument( '--ext', default='.txt', help='dirの場合にキャプションファイルにつける拡張子') - +parser.add_argument( + '--model', + default='wd14-convnextv2.v1', + choices=list(interrogators.keys()), + help='予測に使用するモデル名') args = parser.parse_args() -# 使用するモデル -interrogator = WaifuDiffusionInterrogator( - 'wd14-convnextv2-v2', - repo_id='SmilingWolf/wd-v1-4-convnextv2-tagger-v2', - revision='v2.0' -) +# 使用するinterrogatorを読み出す +interrogator = interrogators[args.model] def image_interrogate(image_path: Path): """ diff --git a/tagger/dbimutils.py b/tagger/dbimutils.py index 93ea643..34a1e06 100644 --- a/tagger/dbimutils.py +++ b/tagger/dbimutils.py @@ -4,6 +4,26 @@ import cv2 import numpy as np from PIL import Image +def fill_transparent(image: Image.Image, color='WHITE'): + image = image.convert('RGBA') + new_image = Image.new('RGBA', image.size, color) + new_image.paste(image, mask=image) + image = new_image.convert('RGB') + return image + +def resize(pic: Image.Image, size: int, keep_ratio=True) -> Image.Image: + if not keep_ratio: + target_size = (size, size) + else: + min_edge = min(pic.size) + target_size = ( + int(pic.size[0] / min_edge * size), + int(pic.size[1] / min_edge * size), + ) + + target_size = (target_size[0] & ~3, target_size[1] & ~3) + + return pic.resize(target_size, resample=Image.Resampling.LANCZOS) def smart_imread(img, flag=cv2.IMREAD_UNCHANGED): if img.endswith(".gif"): diff --git a/tagger/interrogator.py b/tagger/interrogator.py index 974ffe8..dc5dbe2 100644 --- a/tagger/interrogator.py +++ b/tagger/interrogator.py @@ -9,16 +9,15 @@ 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 use_cpu = True -if use_cpu: - tf_device_name = '/cpu:0' -else: - tf_device_name = '/gpu:0' class Interrogator: @staticmethod @@ -187,3 +186,79 @@ class WaifuDiffusionInterrogator(Interrogator): 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 + providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] + if use_cpu: + providers.pop(0) + self.model = InferenceSession(model_path, + providers=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/interrogators.py b/tagger/interrogators.py new file mode 100644 index 0000000..4bf4423 --- /dev/null +++ b/tagger/interrogators.py @@ -0,0 +1,46 @@ +from typing import List, Dict + +from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator, MLDanbooruInterrogator + +interrogators: Dict[str, Interrogator] = { + 'wd14-vit.v1': WaifuDiffusionInterrogator( + 'WD14 ViT v1', + repo_id='SmilingWolf/wd-v1-4-vit-tagger' + ), + 'wd14-vit.v2': WaifuDiffusionInterrogator( + 'WD14 ViT v2', + repo_id='SmilingWolf/wd-v1-4-vit-tagger-v2', + ), + 'wd14-convnext.v1': WaifuDiffusionInterrogator( + 'WD14 ConvNeXT v1', + repo_id='SmilingWolf/wd-v1-4-convnext-tagger' + ), + 'wd14-convnext.v2': WaifuDiffusionInterrogator( + 'WD14 ConvNeXT v2', + repo_id='SmilingWolf/wd-v1-4-convnext-tagger-v2', + ), + 'wd14-convnextv2.v1': WaifuDiffusionInterrogator( + 'WD14 ConvNeXTV2 v1', + # the name is misleading, but it's v1 + repo_id='SmilingWolf/wd-v1-4-convnextv2-tagger-v2', + ), + 'wd14-swinv2-v1': WaifuDiffusionInterrogator( + 'WD14 SwinV2 v1', + # again misleading name + repo_id='SmilingWolf/wd-v1-4-swinv2-tagger-v2', + ), + 'wd-v1-4-moat-tagger.v2': WaifuDiffusionInterrogator( + 'WD14 moat tagger v2', + repo_id='SmilingWolf/wd-v1-4-moat-tagger-v2' + ), + 'mld-caformer.dec-5-97527': MLDanbooruInterrogator( + 'ML-Danbooru Caformer dec-5-97527', + repo_id='deepghs/ml-danbooru-onnx', + model_path='ml_caformer_m36_dec-5-97527.onnx' + ), + 'mld-tresnetd.6-30000': MLDanbooruInterrogator( + 'ML-Danbooru TResNet-D 6-30000', + repo_id='deepghs/ml-danbooru-onnx', + model_path='TResnet-D-FLq_ema_6-30000.onnx' + ), +} \ No newline at end of file