File size: 2,528 Bytes
1f4a582 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import gradio as gr
import huggingface_hub
from PIL import Image
from pathlib import Path
import onnxruntime as rt
import numpy as np
import csv
import spaces
import onnxruntime as rt
e621_model_path = Path(huggingface_hub.snapshot_download('toynya/Z3D-E621-Convnext'))
e621_model_session = rt.InferenceSession(e621_model_path / 'model.onnx', providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
with open(e621_model_path / 'tags-selected.csv', mode='r', encoding='utf-8') as file:
csv_reader = csv.DictReader(file)
e621_model_tags = [row['name'].strip() for row in csv_reader]
def prepare_image_e621(image: Image.Image, target_size: int):
import numpy as np
# Pad image to square
image_shape = image.size
max_dim = max(image_shape)
pad_left = (max_dim - image_shape[0]) // 2
pad_top = (max_dim - image_shape[1]) // 2
padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
padded_image.paste(image, (pad_left, pad_top))
# Resize
if max_dim != target_size:
padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC)
# Convert to numpy array
# Based on the ONNX graph, the model appears to expect inputs in the range of 0-255
image_array = np.asarray(padded_image, dtype=np.float32)
# Convert PIL-native RGB to BGR
image_array = image_array[:, :, ::-1]
return np.expand_dims(image_array, axis=0)
def predict_e621(image: Image.Image):
THRESHOLD = 0.3
image_array = prepare_image_e621(image, 448)
image_array = prepare_image_e621(image, 448)
input_name = 'input_1:0'
output_name = 'predictions_sigmoid'
result = e621_model_session.run([output_name], {input_name: image_array})
result = result[0][0]
scores = {e621_model_tags[i]: result[i] for i in range(len(result))}
predicted_tags = [tag for tag, score in scores.items() if score > THRESHOLD]
tag_string = ', '.join(predicted_tags).replace("_", " ")
return tag_string, scores
DESCRIPTION = """
E621 Tagger (Z3D-E621-Convnext)
- Image => E621 Pony Prompt
- Mod of [fancyfeast's demo](https://huggingface.co/spaces/fancyfeast/Z3D-E621-Convnext-space) for toynya's [Z3D-E621-Convnext](https://huggingface.co/toynya/Z3D-E621-Convnext)
"""
gradio_app = gr.Interface(
predict_e621,
inputs=gr.Image(label="Source", sources=['upload', 'clipboard'], type='pil'),
outputs=[
gr.Textbox(label="Tag String", show_copy_button=True),
gr.Label(label="Tag Predictions", num_top_classes=100),
],
description=DESCRIPTION,
allow_flagging="never",
)
if __name__ == '__main__':
gradio_app.launch()
|