votepurchase/DeepDanbooru

1#!/usr/bin/env python2 3from __future__ import annotations4 5import os6import pathlib7import tarfile8 9import deepdanbooru as dd10import gradio as gr11import huggingface_hub12import numpy as np13import PIL.Image14import tensorflow as tf15 16DESCRIPTION = "# [KichangKim/DeepDanbooru](https://github.com/KichangKim/DeepDanbooru)"17 18 19def load_sample_image_paths() -> list[pathlib.Path]:20    image_dir = pathlib.Path("images")21    if not image_dir.exists():22        path = huggingface_hub.hf_hub_download("public-data/sample-images-TADNE", "images.tar.gz", repo_type="dataset")23        with tarfile.open(path) as f:24            f.extractall()25    return sorted(image_dir.glob("*"))26 27 28def load_model() -> tf.keras.Model:29    path = huggingface_hub.hf_hub_download("public-data/DeepDanbooru", "model-resnet_custom_v3.h5")30    model = tf.keras.models.load_model(path)31    return model32 33 34def load_labels() -> list[str]:35    path = huggingface_hub.hf_hub_download("public-data/DeepDanbooru", "tags.txt")36    with open(path) as f:37        labels = [line.strip() for line in f.readlines()]38    return labels39 40 41model = load_model()42labels = load_labels()43 44 45def predict(image: PIL.Image.Image, score_threshold: float) -> tuple[dict[str, float], dict[str, float], str]:46    _, height, width, _ = model.input_shape47    image = np.asarray(image)48    image = tf.image.resize(image, size=(height, width), method=tf.image.ResizeMethod.AREA, preserve_aspect_ratio=True)49    image = image.numpy()50    image = dd.image.transform_and_pad_image(image, width, height)51    image = image / 255.052    probs = model.predict(image[None, ...])[0]53    probs = probs.astype(float)54 55    indices = np.argsort(probs)[::-1]56    result_all = dict()57    result_threshold = dict()58    for index in indices:59        label = labels[index]60        prob = probs[index]61        result_all[label] = prob62        if prob < score_threshold:63            break64        result_threshold[label] = prob65    result_text = ", ".join(result_all.keys())66    return result_threshold, result_all, result_text67 68 69image_paths = load_sample_image_paths()70examples = [[path.as_posix(), 0.5] for path in image_paths]71 72with gr.Blocks(css="style.css") as demo:73    gr.Markdown(DESCRIPTION)74    with gr.Row():75        with gr.Column():76            image = gr.Image(label="Input", type="pil")77            score_threshold = gr.Slider(label="Score threshold", minimum=0, maximum=1, step=0.05, value=0.5)78            run_button = gr.Button("Run")79        with gr.Column():80            with gr.Tabs():81                with gr.Tab(label="Output"):82                    result = gr.Label(label="Output", show_label=False)83                with gr.Tab(label="JSON"):84                    result_json = gr.JSON(label="JSON output", show_label=False)85                with gr.Tab(label="Text"):86                    result_text = gr.Text(label="Text output", show_label=False, lines=5)87    gr.Examples(88        examples=examples,89        inputs=[image, score_threshold],90        outputs=[result, result_json, result_text],91        fn=predict,92        cache_examples=os.getenv("CACHE_EXAMPLES") == "1",93    )94 95    run_button.click(96        fn=predict,97        inputs=[image, score_threshold],98        outputs=[result, result_json, result_text],99        api_name="predict",100    )101 102if __name__ == "__main__":103    demo.queue(max_size=20).launch()104