Spaces:
Runtime error
Runtime error
File size: 1,876 Bytes
d8240fe 6d58819 ffa3d56 6d58819 f67fcbe 6d58819 ffa3d56 cdab4c3 ffa3d56 6d58819 f67fcbe 6d58819 9f52a4a 6d58819 d8240fe 6d58819 d8240fe 6d58819 |
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 |
import os
from functools import lru_cache
import gradio as gr
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
from imgutils.data import load_image
from imgutils.utils import open_onnx_model
import gradio as gr
from huggingface_hub import hf_hub_download
from onnx import hub
import onnxruntime as ort
_MODELS = [
('content_moderation.onnx', 224),
]
_MODEL_NAMES = [name for name, _ in _MODELS]
_DEFAULT_MODEL_NAME = _MODEL_NAMES[0]
_MODEL_TO_SIZE = dict(_MODELS)
model = hub.load(hf_hub_download(
'tanlocc/Out_of_Universe',
"content_moderation.onnx"
))
print(model)
@lru_cache()
def _onnx_model(name):
return open_onnx_model(hf_hub_download(
'tanlocc/Out_of_Universe',
f'{name}'
))
def _image_preprocess(image, size: int = 224) -> np.ndarray:
image = load_image(image, mode='RGB').resize((size, size), Image.NEAREST)
return (np.array(image) / 255.0)[None, ...]
_LABELS = ['drawings', 'hentai', 'neutral', 'porn', 'sexy']
def predict(image, model_name):
input_ = _image_preprocess(image, _MODEL_TO_SIZE[model_name]).astype(np.float32)
output_, = _onnx_model(model_name).run()
return dict(zip(_LABELS, map(float, output_[0])))
if __name__ == '__main__':
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
gr_input_image = gr.Image(type='pil', label='Original Image')
gr_model = gr.Dropdown(_MODEL_NAMES, value=_DEFAULT_MODEL_NAME, label='Model')
gr_btn_submit = gr.Button(value='Tagging', variant='primary')
with gr.Column():
gr_ratings = gr.Label(label='Ratings')
gr_btn_submit.click(
predict,
inputs=[gr_input_image, gr_model],
outputs=[gr_ratings],
)
demo.queue(os.cpu_count()).launch()
|