Load model directly from hub

app.py

@@ -3,7 +3,7 @@ import requests
 import numpy as np
 import streamlit as st
 import tensorflow as tf
-from huggingface_hub import snapshot_download
+from huggingface_hub import from_pretrained_keras
 from tensorflow.keras.preprocessing.image import img_to_array
 from tensorflow.keras.models import load_model
 from PIL import Image
@@ -18,10 +18,10 @@ def make_prediction(url, model, race_names):
 
     predictions = model.predict(img_array)
 
-    top4 = predictions.argsort()[0, -1:-5:-1]
+    top5 = predictions.argsort()[0, -1:-6:-1]
 
     breakdown = []
-    for race, acc in zip(np.array(race_names)[top4], predictions[0, top4]):
+    for race, acc in zip(np.array(race_names)[top5], predictions[0, top5]):
         breakdown.append(f'{race} at {acc:.2%}')
     return breakdown
 
@@ -31,8 +31,7 @@ def main():
     race_file = open('race_names.json', 'r')
     race_names = json.load(race_file)
     # Load model
-    download_folder = snapshot_download(repo_id='JbIPS/DogRace')
-    model = load_model(download_folder)
+    model = from_pretrained_keras('JbIPS/DogRace')
 
     st.set_page_config("Who let's the dogs out")
     st.title('Quelle est ta race de chien totem ?')