Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,11 +1,11 @@
 import streamlit as st
+import tensorflow as tf
 import numpy as np
 from PIL import Image
-import onnxruntime as ort
 
 # ====== MODEL SETTINGS ======
-MODEL_PATH = "cnn_largefish_model.onnx"
-IMG_SIZE = 64  # jouw trainingsgrootte
+MODEL_PATH = "cnn_largefish_model.h5"  # op HuggingFace gewoon in de root neerzetten
+IMG_SIZE = 64  # jouw resize in Kaggle
 
 CLASS_NAMES = [
     'House Mackerel',
@@ -19,37 +19,55 @@ CLASS_NAMES = [
     'Red Sea Bream'
 ]
 
+
+# ====== FUNCTIES ======
 @st.cache_resource
-def load_session():
-    session = ort.InferenceSession(
-        MODEL_PATH,
-        providers=["CPUExecutionProvider"]
-    )
-    input_name = session.get_inputs()[0].name
-    return session, input_name
+def load_model():
+    """Laad het Keras-model één keer en cache het."""
+    model = tf.keras.models.load_model(MODEL_PATH)
+    return model
+
 
 def preprocess_image(image: Image.Image) -> np.ndarray:
+    """Resize + normaliseer afbeelding naar hetzelfde formaat als training."""
     image = image.convert("RGB")
     image = image.resize((IMG_SIZE, IMG_SIZE))
     arr = np.array(image).astype("float32") / 255.0
-    arr = np.expand_dims(arr, axis=0)  
+    arr = np.expand_dims(arr, axis=0)  # shape: (1, 64, 64, 3)
     return arr
 
-def predict(image: Image.Image):
-    session, input_name = load_session()
+
+def predict_image(model, image: Image.Image):
+    """Voorspel klasse + probabilities voor één afbeelding."""
     x = preprocess_image(image)
-    preds = session.run(None, {input_name: x})[0][0]
+    preds = model.predict(x)[0]  # shape: (9,)
     pred_idx = int(np.argmax(preds))
     pred_class = CLASS_NAMES[pred_idx]
     pred_conf = float(preds[pred_idx])
     return pred_class, pred_conf, preds
 
-st.set_page_config(page_title="Fish Classifier", page_icon="🐟")
 
-st.title("🐟 Large-Scale Fish Classifier (ONNX)")
-st.write("Upload een afbeelding en het model voorspelt de soort.")
+# ====== STREAMLIT UI ======
+st.set_page_config(
+    page_title="Fish Classifier",
+    page_icon="🐟",
+)
 
-uploaded_file = st.file_uploader("Upload een afbeelding", type=["jpg", "jpeg", "png"])
+st.title("🐟 Large-Scale Fish Classifier")
+st.write(
+    """
+Upload een afbeelding van een vis uit de dataset 
+en het model voorspelt de soort.
+"""
+)
+
+# Model alvast laden (toont spinner)
+with st.spinner("Model laden..."):
+    model = load_model()
+
+uploaded_file = st.file_uploader(
+    "Upload een afbeelding", type=["jpg", "jpeg", "png"]
+)
 
 if uploaded_file is not None:
     image = Image.open(uploaded_file)
@@ -57,12 +75,15 @@ if uploaded_file is not None:
 
     if st.button("Classify"):
         with st.spinner("Bezig met voorspellen..."):
-            pred_class, pred_conf, preds = predict(image)
+            pred_class, pred_conf, preds = predict_image(model, image)
 
         st.subheader("Voorspelling")
         st.write(f"**{pred_class}** met **{pred_conf:.2%}** zekerheid.")
 
+        # Probabilities plotten als bar chart
+        prob_dict = {CLASS_NAMES[i]: float(preds[i]) for i in range(len(CLASS_NAMES))}
         st.subheader("Class probabilities")
-        st.bar_chart({CLASS_NAMES[i]: float(preds[i]) for i in range(len(CLASS_NAMES))})
+        st.bar_chart(prob_dict)
+
 else:
-    st.info("➡️ Upload eerst een afbeelding.")
+    st.info("➡️ Upload eerst een afbeelding (jpg/jpeg/png).")