Update app.py

app.py

@@ -2,53 +2,81 @@ import os
 import numpy as np
 import streamlit as st
 from PIL import Image
-import tensorflow as tf
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.image import img_to_array
 
-# Set page config
-st.set_page_config(page_title="Fish Classifier", page_icon="🐟")
+# CONFIG
+MODEL_PATH = "custom_cnn_last4_finetuned (1).h5"
+IMG_SIZE = (256, 256)
 
-st.title("🐟 Fish Species Classifier")
+CLASS_NAMES = [
+    "animal fish",
+    "animal fish bass",
+    "fish sea_food black_sea_sprat",
+    "fish sea_food gilt_head_bream",
+    "fish sea_food hourse_mackerel",
+    "fish sea_food red_mullet",
+    "fish sea_food red_sea_bream",
+    "fish sea_food sea_bass",
+    "fish sea_food shrimp",
+    "fish sea_food striped_red_mullet",
+    "fish sea_food trout"
+]
 
-# List all files to find the model
-st.write("Searching for model file...")
+st.set_page_config(page_title="Custom CNN Fish Classifier", layout="centered")
+st.title("🐟 Fish Classifier")
 
-# Look for model in common locations
-possible_paths = [
-    "EfficientNetB0_head_finetuned (1).keras",
-    "EfficientNetB0_head_finetuned.keras",
-    "model.keras",
-    "fish_model.keras",
-    "EfficientNetB0.keras",
-]
+# LOAD MODEL
+@st.cache_resource
+def load_cnn_model():
+    try:
+        model = load_model(MODEL_PATH, compile=False)
+        return model
+    except Exception as e:
+        st.error(f"Model loading failed:\n{e}")
+        st.info("""
+        **Upload your model file to this Space:**
+        File must be named: `custom_cnn_last4_finetuned (1).h5`
+        """)
+        return None
 
-found_model = None
-for path in possible_paths:
-    if os.path.exists(path):
-        st.success(f"✅ Found model at: {path}")
-        found_model = path
-        break
+model = load_cnn_model()
 
-if not found_model:
-    st.error("❌ Model file not found!")
-    st.write("Please upload your model file to the Space.")
+if model is None:
+    # Show what files exist
+    if os.path.exists("."):
+        st.write("Files in this space:")
+        for f in os.listdir("."):
+            st.write(f"- {f}")
     st.stop()
 
-# Try to load the model
-try:
-    model = tf.keras.models.load_model(found_model, compile=False)
-    st.success("✅ Model loaded successfully!")
-except Exception as e:
-    st.error(f"❌ Failed to load model: {str(e)[:200]}")
-    st.stop()
+# PREPROCESS IMAGE
+def prepare_image(pil_img):
+    pil_img = pil_img.convert("RGB")
+    pil_img = pil_img.resize((IMG_SIZE[1], IMG_SIZE[0]))
+
+    arr = img_to_array(pil_img)
+    arr = arr / 255.0  # Normalize to 0-1
+    arr = np.expand_dims(arr, axis=0)
 
-# Simple interface
-uploaded = st.file_uploader("Upload fish image", type=["jpg", "png", "jpeg"])
+    return arr
+
+# PREDICT
+def predict_top1(img):
+    x = prepare_image(img)
+    preds = model.predict(x, verbose=0)[0]
+    top_index = np.argmax(preds)
+
+    return CLASS_NAMES[top_index], float(preds[top_index])
+
+# UI
+uploaded = st.file_uploader("Upload fish image", type=["jpg", "jpeg", "png"])
 
 if uploaded:
     img = Image.open(uploaded)
-    st.image(img, caption="Uploaded Image")
-    
-    if st.button("Classify"):
-        st.write("Processing...")
-        # Add your prediction code here
-        st.success("Classification complete!")
+    st.image(img, caption="Uploaded Image", use_container_width=True)
+
+    if st.button("Predict"):
+        label, prob = predict_top1(img)
+        st.markdown(f"## 🎯 Prediction: **{label}**")
+        st.markdown(f"### Confidence: **{prob*100:.2f}%**")