Spaces:
Running
Running
File size: 3,047 Bytes
fb59cb8 e9c5f95 fb59cb8 01b28b7 fb59cb8 e9c5f95 01b28b7 e9c5f95 fb59cb8 e9c5f95 01b28b7 e9c5f95 fb59cb8 e9c5f95 01b28b7 e9c5f95 01b28b7 fb59cb8 18988d2 fb59cb8 5d94de1 fb59cb8 01b28b7 |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
#!/usr/bin/env python
from __future__ import annotations
import os
import pathlib
import tarfile
import deepdanbooru as dd
import gradio as gr
import huggingface_hub
import numpy as np
import PIL.Image
import tensorflow as tf
DESCRIPTION = '# [KichangKim/DeepDanbooru](https://github.com/KichangKim/DeepDanbooru)'
def load_sample_image_paths() -> list[pathlib.Path]:
image_dir = pathlib.Path('images')
if not image_dir.exists():
path = huggingface_hub.hf_hub_download(
'public-data/sample-images-TADNE',
'images.tar.gz',
repo_type='dataset')
with tarfile.open(path) as f:
f.extractall()
return sorted(image_dir.glob('*'))
def load_model() -> tf.keras.Model:
path = huggingface_hub.hf_hub_download('public-data/DeepDanbooru',
'model-resnet_custom_v3.h5')
model = tf.keras.models.load_model(path)
return model
def load_labels() -> list[str]:
path = huggingface_hub.hf_hub_download('public-data/DeepDanbooru',
'tags.txt')
with open(path) as f:
labels = [line.strip() for line in f.readlines()]
return labels
model = load_model()
labels = load_labels()
def predict(image: PIL.Image.Image,
score_threshold: float) -> dict[str, float]:
_, height, width, _ = model.input_shape
image = np.asarray(image)
image = tf.image.resize(image,
size=(height, width),
method=tf.image.ResizeMethod.AREA,
preserve_aspect_ratio=True)
image = image.numpy()
image = dd.image.transform_and_pad_image(image, width, height)
image = image / 255.
probs = model.predict(image[None, ...])[0]
probs = probs.astype(float)
res = dict()
for prob, label in zip(probs.tolist(), labels):
if prob < score_threshold:
continue
res[label] = prob
return res
image_paths = load_sample_image_paths()
examples = [[path.as_posix(), 0.5] for path in image_paths]
with gr.Blocks(css='style.css') as demo:
gr.Markdown(DESCRIPTION)
with gr.Row():
with gr.Column():
image = gr.Image(label='Input', type='pil')
score_threshold = gr.Slider(label='Score threshold',
minimum=0,
maximum=1,
step=0.05,
value=0.5)
run_button = gr.Button('Run')
with gr.Column():
result = gr.Label(label='Output')
gr.Examples(examples=examples,
inputs=[image, score_threshold],
outputs=result,
fn=predict,
cache_examples=os.getenv('CACHE_EXAMPLES') == '1')
run_button.click(fn=predict,
inputs=[image, score_threshold],
outputs=result,
api_name='predict')
demo.queue().launch()
|