Spaces:
Paused
Paused
#!/usr/bin/env python | |
from __future__ import annotations | |
import argparse | |
import functools | |
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 | |
TOKEN = os.environ['TOKEN'] | |
MODEL_REPO = 'hysts/DeepDanbooru' | |
MODEL_FILENAME = 'model-resnet_custom_v3.h5' | |
LABEL_FILENAME = 'tags.txt' | |
def parse_args() -> argparse.Namespace: | |
parser = argparse.ArgumentParser() | |
parser.add_argument('--score-slider-step', type=float, default=0.05) | |
parser.add_argument('--score-threshold', type=float, default=0.5) | |
parser.add_argument('--theme', type=str, default='dark-grass') | |
parser.add_argument('--live', action='store_true') | |
parser.add_argument('--share', action='store_true') | |
parser.add_argument('--port', type=int) | |
parser.add_argument('--disable-queue', | |
dest='enable_queue', | |
action='store_false') | |
parser.add_argument('--allow-flagging', type=str, default='never') | |
parser.add_argument('--allow-screenshot', action='store_true') | |
return parser.parse_args() | |
def load_sample_image_paths() -> list[pathlib.Path]: | |
image_dir = pathlib.Path('images') | |
if not image_dir.exists(): | |
dataset_repo = 'hysts/sample-images-TADNE' | |
path = huggingface_hub.hf_hub_download(dataset_repo, | |
'images.tar.gz', | |
repo_type='dataset', | |
use_auth_token=TOKEN) | |
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(MODEL_REPO, | |
MODEL_FILENAME, | |
use_auth_token=TOKEN) | |
model = tf.keras.models.load_model(path) | |
return model | |
def load_labels() -> list[str]: | |
path = huggingface_hub.hf_hub_download(MODEL_REPO, | |
LABEL_FILENAME, | |
use_auth_token=TOKEN) | |
with open(path) as f: | |
labels = [line.strip() for line in f.readlines()] | |
return labels | |
def predict(image: PIL.Image.Image, score_threshold: float, | |
model: tf.keras.Model, labels: list[str]) -> 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, labels): | |
if prob < score_threshold: | |
continue | |
res[label] = prob | |
return res | |
def main(): | |
gr.close_all() | |
args = parse_args() | |
image_paths = load_sample_image_paths() | |
examples = [[path.as_posix(), args.score_threshold] | |
for path in image_paths] | |
model = load_model() | |
labels = load_labels() | |
func = functools.partial(predict, model=model, labels=labels) | |
func = functools.update_wrapper(func, predict) | |
repo_url = 'https://github.com/KichangKim/DeepDanbooru' | |
title = 'KichangKim/DeepDanbooru' | |
description = f'A demo for {repo_url}' | |
article = None | |
gr.Interface( | |
func, | |
[ | |
gr.inputs.Image(type='pil', label='Input'), | |
gr.inputs.Slider(0, | |
1, | |
step=args.score_slider_step, | |
default=args.score_threshold, | |
label='Score Threshold'), | |
], | |
gr.outputs.Label(label='Output'), | |
theme=args.theme, | |
title=title, | |
description=description, | |
article=article, | |
examples=examples, | |
allow_screenshot=args.allow_screenshot, | |
allow_flagging=args.allow_flagging, | |
live=args.live, | |
).launch( | |
enable_queue=args.enable_queue, | |
server_port=args.port, | |
share=args.share, | |
) | |
if __name__ == '__main__': | |
main() | |