Refactor.
This commit is contained in:
corkborg
@@ -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)
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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)}")
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
@@ -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'
|
||||
|
||||
Reference in New Issue
Block a user