push keras file

src/prediction.py

@@ -2,6 +2,8 @@ import numpy as np
 import streamlit as st
 import tensorflow as tf
 from PIL import Image
+from huggingface_hub import hf_hub_download
+
 
 class_names = [
  'Tomato_Bacterial_spot',
@@ -17,8 +19,14 @@ class_names = [
 ]
 
 @st.cache_resource
-def load_my_model():
-    return tf.keras.models.load_model("leaf_detection_model.keras")
+def load_model_from_hub():
+    model_path = hf_hub_download(
+        repo_id="Heizsenberg/leaf-classification-tomato",
+        filename="leaf_detection_model.keras"
+    )
+    model = tf.keras.models.load_model(model_path)
+    return model
+
 
 
 def predict_image(model, img):
@@ -38,7 +46,7 @@ def predict_image(model, img):
     return img, predicted_label, confidence
 
 def run():
-    classifier_model = load_my_model()
+    classifier_model = load_model_from_hub()
     
     st.write("upload a tomato leaf image to be predicted")