Update

.gitmodules

@@ -1,3 +0,0 @@
-[submodule "DeepDanbooru"]
-	path = DeepDanbooru
-	url = https://github.com/KichangKim/DeepDanbooru

DeepDanbooru

@@ -1 +0,0 @@
-Subproject commit 92ba0b56be5eed0037e3f067bb9867f5ac691647

app.py

@@ -7,19 +7,12 @@ import functools
 import os
 import pathlib
 import subprocess
-import sys
-import urllib
-import zipfile
-from typing import Callable
+import tarfile
 
 # workaround for https://github.com/gradio-app/gradio/issues/483
 command = 'pip install -U gradio==2.7.0'
 subprocess.call(command.split())
 
-command = 'pip install -r DeepDanbooru/requirements.txt'
-subprocess.call(command.split())
-sys.path.insert(0, 'DeepDanbooru')
-
 import deepdanbooru as dd
 import gradio as gr
 import huggingface_hub
@@ -29,9 +22,9 @@ import tensorflow as tf
 
 TOKEN = os.environ['TOKEN']
 
-ZIP_PATH = 'data.zip'
-TAG_PATH = 'tags.txt'
-MODEL_PATH = 'model-resnet_custom_v3.h5'
+MODEL_REPO = 'hysts/DeepDanbooru'
+MODEL_FILENAME = 'model-resnet_custom_v3.h5'
+LABEL_FILENAME = 'tags.txt'
 
 
 def parse_args() -> argparse.Namespace:
@@ -50,33 +43,38 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def download_sample_images() -> list[pathlib.Path]:
-    image_dir = pathlib.Path('samples')
-    image_dir.mkdir(exist_ok=True)
-
-    dataset_repo = 'hysts/sample-images-TADNE'
-    n_images = 36
-    paths = []
-    for index in range(n_images):
+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,
-                                               f'{index:02d}.jpg',
+                                               'images.tar.gz',
                                                repo_type='dataset',
-                                               cache_dir=image_dir.as_posix(),
                                                use_auth_token=TOKEN)
-        paths.append(pathlib.Path(path))
-    return paths
+        with tarfile.open(path) as f:
+            f.extractall()
+    return sorted(image_dir.glob('*'))
 
 
-def download_model_data() -> None:
-    url = 'https://github.com/KichangKim/DeepDanbooru/releases/download/v3-20200915-sgd-e30/deepdanbooru-v3-20200915-sgd-e30.zip'
-    urllib.request.urlretrieve(url, ZIP_PATH)
-    with zipfile.ZipFile(ZIP_PATH) as f:
-        f.extract(TAG_PATH)
-        f.extract(MODEL_PATH)
+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 predict(image: PIL.Image.Image, score_threshold: float, model,
-            labels: list[str]) -> dict[str, float]:
+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,
@@ -101,18 +99,12 @@ def main():
 
     args = parse_args()
 
-    image_paths = download_sample_images()
+    image_paths = load_sample_image_paths()
     examples = [[path.as_posix(), args.score_threshold]
                 for path in image_paths]
 
-    zip_path = pathlib.Path(ZIP_PATH)
-    if not zip_path.exists():
-        download_model_data()
-
-    model = tf.keras.models.load_model(MODEL_PATH)
-
-    with open(TAG_PATH) as f:
-        labels = [line.strip() for line in f.readlines()]
+    model = load_model()
+    labels = load_labels()
 
     func = functools.partial(predict, model=model, labels=labels)
     func = functools.update_wrapper(func, predict)

requirements.txt

@@ -1,2 +1,3 @@
 pillow>=9.0.0
 tensorflow>=2.7.0
+git+https://github.com/KichangKim/DeepDanbooru@v3-20200915-sgd-e30#egg=deepdanbooru