258 lines
7.2 KiB
Python
258 lines
7.2 KiB
Python
import os
|
|
import pandas as pd
|
|
import numpy as np
|
|
|
|
from typing import Tuple, List, Dict
|
|
from io import BytesIO
|
|
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: List[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 = ['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_gpu(self) -> None:
|
|
self.providers = ['CUDAExecutionProvider']
|
|
|
|
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()
|