download models via hf_hub_download

app.py

@@ -14,8 +14,9 @@ from subprocess import run
 import json
 
 import gradio as gr
+from huggingface_hub import hf_hub_download
 from yolov5 import detect
-from utils import get_models, get_cfg, get_embeddings, get_comp_embeddings, get_test_embedding
+from utils import get_model, get_cfg, get_embeddings, get_comp_embeddings, get_test_embedding
 
 
 # YOLOv5 parameters
@@ -25,8 +26,8 @@ score_thr = 0.025
 iou_thr = 0.6
 max_det = 1
 working = Path(os.getcwd())
-modelbox = working / 'models'
-checkpoint_files = [modelbox / f'yolov5_l6_{yolo_input_size}_fold{x}.pt' for x in versions]
+modelbox = "yellowdolphin/happywhale-models"
+checkpoint_files = [hf_hub_download(modelbox, f'yolov5_l6_{yolo_input_size}_fold{x}.pt') for x in versions]
 image_root = working / 'images'
 
 
@@ -38,10 +39,11 @@ normalize_similarity = None  # test-train, None
 gamma = 0.4
 threshold = 0.09951 if (normalize_similarity == 'test-train') else 0.6 # 0.381
 knn = 300
-emb_path = '/kaggle/input/happywhale-embeddings'
-rst_path = '/kaggle/input/happywhale-models'
-rst_files = sorted(glob(f'{rst_path}/*.h5'))
-n_models = len(rst_files)
+rst_names = 'convnext_base_384_in22ft1k_colab220 efnv1b7_colab216 hub_efnv2xl_v73'.split()
+cfg_files = [hf_hub_download(modelbox, f'{x}_config.json') for x in rst_names]
+emb_files = [hf_hub_download(modelbox, f'{x}_emb.npz') for x in rst_names]
+rst_files = [hf_hub_download(modelbox, f'{x}.h5') for x in rst_names]
+n_models = len(rst_names)
 
 
 def fast_yolo_crop(image):
@@ -74,14 +76,13 @@ def fast_yolo_crop(image):
 
 
 # Preload embeddings for known individuals
-comp_embeddings = get_comp_embeddings(rst_files)
+comp_embeddings = get_comp_embeddings(emb_files)
 
 # Preload embedding models, input sizes
 K.clear_session()
 embed_models, sizes = [], []
-for rst_file in rst_files:
-    cfg = get_cfg(rst_file)
-    npz_file = Path(rst_file.replace('.h5', '_emb.npz')).name
+for cfg_file, rst_file, npz_file in zip (cfg_files, rst_files, emb_files):
+    cfg = get_cfg(cfg_file)
     assert cfg.FOLD_TO_RUN == use_fold[npz_file]
     cfg.pretrained = None  # avoid weight downloads
     if isinstance(cfg.IMAGE_SIZE, int):

utils.py

@@ -19,8 +19,8 @@ class DotDict(dict):
     __delattr__ = dict.__delitem__
 
     
-def get_cfg(rst_file):
-    json_file = str(rst_file).replace('.h5', '_config.json')
+def get_cfg(json_file):
+    json_file = str(json_file)
     config_dict = json.load(open(json_file))
     return DotDict(config_dict)
 
@@ -77,15 +77,14 @@ use_fold = {
 }
 
 
-def get_comp_embeddings(rst_files):
+def get_comp_embeddings(emb_files):
     "Load embeddings for competition images [n_images, embedding_size]"
 
     comp_embeddings = []
 
-    for rst_file in rst_files:
+    for npz_file in emb_files:
         # Get embeddings for all competition images
-        npz_file = Path(rst_file.replace('.h5', '_emb.npz')).name
-        d = np.load(str(Path(emb_path) / npz_file))
+        d = np.load(str(npz_file))
         comp_train_emb = d['train']
         comp_test_emb = d['test']