Update

app.py

@@ -8,6 +8,7 @@ import os
 import pathlib
 import subprocess
 import sys
+import tarfile
 
 # workaround for https://github.com/gradio-app/gradio/issues/483
 command = 'pip install -U gradio==2.7.0'
@@ -26,8 +27,8 @@ from _util.twodee_v0 import I as ImageWrapper
 TOKEN = os.environ['TOKEN']
 
 MODEL_REPO = 'hysts/bizarre-pose-estimator-models'
-MODEL_PATH = 'tagger.pth'
-LABEL_PATH = 'tags.txt'
+MODEL_FILENAME = 'tagger.pth'
+LABEL_FILENAME = 'tags.txt'
 
 
 def parse_args() -> argparse.Namespace:
@@ -47,21 +48,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 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 = torchvision.models.resnet50(num_classes=1062)
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    return model
+
+
+def load_labels() -> list[str]:
+    label_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                 LABEL_FILENAME,
+                                                 use_auth_token=TOKEN)
+    with open(label_path) as f:
+        labels = [line.strip() for line in f.readlines()]
+    return labels
 
 
 @torch.inference_mode()
@@ -84,36 +102,13 @@ def predict(image: PIL.Image.Image, score_threshold: float,
     return res
 
 
-def load_model(device: torch.device) -> torch.nn.Module:
-    model_path = huggingface_hub.hf_hub_download(MODEL_REPO,
-                                                 MODEL_PATH,
-                                                 use_auth_token=TOKEN)
-    state_dict = torch.load(model_path)
-
-    model = torchvision.models.resnet50(num_classes=1062)
-    model.load_state_dict(state_dict)
-    model.to(device)
-    model.eval()
-
-    return model
-
-
-def load_labels() -> list[str]:
-    label_path = huggingface_hub.hf_hub_download(MODEL_REPO,
-                                                 LABEL_PATH,
-                                                 use_auth_token=TOKEN)
-    with open(label_path) as f:
-        labels = [line.strip() for line in f.readlines()]
-    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]