Update

app.py

@@ -8,6 +8,7 @@ import json
 import os
 import pathlib
 import subprocess
+import tarfile
 from typing import Callable
 
 # workaround for https://github.com/gradio-app/gradio/issues/483
@@ -22,6 +23,10 @@ import torchvision.transforms as T
 
 TOKEN = os.environ['TOKEN']
 
+MODEL_REPO = 'hysts/danbooru-pretrained'
+MODEL_FILENAME = 'resnet50-13306192.pth'
+LABEL_FILENAME = 'class_names_6000.json'
+
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
@@ -40,21 +45,40 @@ 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 load_model(device: torch.device) -> torch.nn.Module:
+    path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                           MODEL_FILENAME,
+                                           use_auth_token=TOKEN)
+    state_dict = torch.load(path)
+    model = torch.hub.load('RF5/danbooru-pretrained',
+                           'resnet50',
+                           pretrained=False)
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    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 = json.load(f)
+    return labels
 
 
 @torch.inference_mode()
@@ -75,40 +99,18 @@ def predict(image: PIL.Image.Image, score_threshold: float,
     return res
 
 
-def load_labels() -> list[str]:
-    label_path = pathlib.Path('class_names_6000.json')
-    label_url = 'https://raw.githubusercontent.com/RF5/danbooru-pretrained/master/config/class_names_6000.json'
-    if not label_path.exists():
-        torch.hub.download_url_to_file(label_url, label_path.as_posix())
-    with open(label_path) as f:
-        labels = json.load(f)
-    return labels
-
-
 def main():
     gr.close_all()
 
     args = parse_args()
     device = torch.device(args.device)
 
-    image_paths = download_sample_images()
+    image_paths = load_sample_image_paths()
     examples = [[path.as_posix(), args.score_threshold]
                 for path in image_paths]
 
-    if device.type == 'cpu':
-        model_path = pathlib.Path('resnet50-13306192.pth')
-        model_url = 'https://github.com/RF5/danbooru-pretrained/releases/download/v0.1/resnet50-13306192.pth'
-        if not model_path.exists():
-            torch.hub.download_url_to_file(model_url, model_path.as_posix())
-        model = torch.hub.load('RF5/danbooru-pretrained',
-                               'resnet50',
-                               pretrained=False)
-        state_dict = torch.load(model_path, map_location=device)
-        model.load_state_dict(state_dict)
-    else:
-        model = torch.hub.load('RF5/danbooru-pretrained', 'resnet50')
-    model.to(device)
-    model.eval()
+    model = load_model(device)
+    labels = load_labels()
 
     transform = T.Compose([
         T.Resize(360),
@@ -117,8 +119,6 @@ def main():
                     std=[0.2970, 0.3017, 0.2979]),
     ])
 
-    labels = load_labels()
-
     func = functools.partial(predict,
                              transform=transform,
                              device=device,