Update

README.md

@@ -4,7 +4,7 @@ emoji: 🏃
 colorFrom: gray
 colorTo: purple
 sdk: gradio
-sdk_version: 3.6
+sdk_version: 3.32
 app_file: app.py
 pinned: false
 ---
@@ -27,7 +27,7 @@ Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gr
 Can be either `gradio`, `streamlit`, or `static`
 
 `sdk_version` : _string_
-Only applicable for `streamlit` SDK.  
+Only applicable for `streamlit` SDK.
 See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
 
 `app_file`: _string_

app.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-import argparse
-import functools
 import os
 import pathlib
 import tarfile
@@ -15,56 +13,42 @@ import numpy as np
 import PIL.Image
 import tensorflow as tf
 
-TITLE = 'KichangKim/DeepDanbooru'
-DESCRIPTION = 'This is an unofficial demo for https://github.com/KichangKim/DeepDanbooru.'
-ARTICLE = '<center><img src="https://visitor-badge.glitch.me/badge?page_id=hysts.deepdanbooru" alt="visitor badge"/></center>'
-
-HF_TOKEN = os.environ['HF_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('--share', action='store_true')
-    return parser.parse_args()
+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():
-        dataset_repo = 'hysts/sample-images-TADNE'
-        path = huggingface_hub.hf_hub_download(dataset_repo,
-                                               'images.tar.gz',
-                                               repo_type='dataset',
-                                               use_auth_token=HF_TOKEN)
+        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(MODEL_REPO,
-                                           MODEL_FILENAME,
-                                           use_auth_token=HF_TOKEN)
+    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(MODEL_REPO,
-                                           LABEL_FILENAME,
-                                           use_auth_token=HF_TOKEN)
+    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
 
 
-def predict(image: PIL.Image.Image, score_threshold: float,
-            model: tf.keras.Model, labels: list[str]) -> dict[str, float]:
+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,
@@ -84,39 +68,30 @@ def predict(image: PIL.Image.Image, score_threshold: float,
     return res
 
 
-def main():
-    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)
-
-    gr.Interface(
-        func,
-        [
-            gr.Image(type='pil', label='Input'),
-            gr.Slider(0,
-                      1,
-                      step=args.score_slider_step,
-                      value=args.score_threshold,
-                      label='Score Threshold'),
-        ],
-        gr.Label(label='Output'),
-        examples=examples,
-        title=TITLE,
-        description=DESCRIPTION,
-        article=ARTICLE,
-        allow_flagging='never',
-    ).launch(
-        enable_queue=True,
-        share=args.share,
-    )
-
-
-if __name__ == '__main__':
-    main()
+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()

style.css

@@ -0,0 +1,3 @@
+h1 {
+  text-align: center;
+}