Update streamlit_app.py

streamlit_app.py

@@ -1,74 +1,70 @@
 import streamlit as st
-import numpy as np
+from huggingface_hub import hf_hub_download
 from tensorflow.keras.models import load_model
 from tensorflow.keras.preprocessing import image
+import numpy as np
 from PIL import Image
-import matplotlib.pyplot as plt
-import io
 
-# Load model
+# Load model from Hugging Face Model Hub
 @st.cache_resource
-def load_model_safe():
-    model_path = "fish_freshness_model_retrained_final.keras"
-    try:
-        model = load_model(model_path)
-        return model
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-        return None
+def load_model_from_hf():
+    model_path = hf_hub_download(
+        repo_id="1Codephoenix/fish-freshness-model",
+        filename="fish_freshness_model_retrained_final.keras"
+    )
+    return load_model(model_path)
 
-model = load_model_safe()
+model = load_model_from_hf()
 
-# Define class names and messages
+# Class labels and custom messages
 class_names = ['Fresh', 'Moderately Fresh', 'Spoiled']
 custom_messages = {
     'Fresh': (
         "✅ **Fresh Fish Detected**\n"
-        "- Estimated Age: Less than 1 day old\n"
-        "- Bright eyes, red gills, firm flesh\n"
-        "- Safe for raw or cooked consumption."
+        "- Age: Less than 1 day old\n"
+        "- Characteristics: Bright eyes, red gills, firm flesh"
     ),
     'Moderately Fresh': (
-        "⚠️ **Moderately Fresh Fish Detected**\n"
-        "- Estimated Age: 2–3 days\n"
-        "- Slight odor, softening flesh\n"
-        "- Cook thoroughly before eating."
+        "⚠️ **Moderately Fresh**\n"
+        "- Age: 2–3 days\n"
+        "- Characteristics: Slightly dull, odor begins, flesh softens"
     ),
     'Spoiled': (
-        "🚫 **Spoiled or Unsafe Fish Detected**\n"
-        "- Estimated Age: 4+ days\n"
-        "- Dull eyes, strong odor, possible chemical treatment\n"
-        "- Not safe for consumption. Discard immediately."
+        "🚫 **Spoiled Fish**\n"
+        "- Age: 4+ days or treated\n"
+        "- Characteristics: Cloudy eyes, odor, unsafe to eat"
     )
 }
 
 # Streamlit UI
 st.title("🐟 Fish Freshness Classifier")
-st.subheader("Upload an image of a fish to analyze its freshness")
+st.subheader("Upload an image of a fish to predict its freshness level")
 
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Upload Fish Image", type=["jpg", "jpeg", "png"])
 
-if uploaded_file and model:
+if uploaded_file:
     try:
+        # Load and preprocess the image
         img = Image.open(uploaded_file).convert("RGB")
-        img_resized = img.resize((224, 224))  # Match model input size
+        img_resized = img.resize((224, 224))  # Match model input
         img_array = image.img_to_array(img_resized)
-        img_array = np.expand_dims(img_array / 255.0, axis=0)
+        img_array = np.expand_dims(img_array / 255.0, axis=0)  # Normalize
 
+        # Predict
         prediction = model.predict(img_array)
         predicted_index = np.argmax(prediction)
-        confidence = float(np.max(prediction))
         predicted_class = class_names[predicted_index]
+        confidence = prediction[0][predicted_index]
 
-        # Display results
-        st.image(img, caption=f"Uploaded Image", use_column_width=True)
-        st.markdown(f"### 🎯 Prediction: **{predicted_class}** ({confidence*100:.2f}% confidence)")
+        # Show image and result
+        st.image(img, caption="Uploaded Fish Image", use_column_width=True)
+        st.markdown(f"### 🎯 Prediction: **{predicted_class}** ({confidence * 100:.2f}% confidence)")
         st.markdown(custom_messages[predicted_class])
 
-        # Show confidence for all classes
-        st.subheader("📊 Class Probabilities:")
+        # Show all class probabilities
+        st.subheader("📊 Confidence Scores")
         for i, class_name in enumerate(class_names):
             st.write(f"- {class_name}: {prediction[0][i]*100:.2f}%")
 
     except Exception as e:
-        st.error(f"Error processing image: {e}")
+        st.error(f"⚠️ Error processing image: {e}")