Implement camie-tagger.
This commit is contained in:
corkborg
@@ -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
|
||||
|
||||
@@ -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'
|
||||
),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user