Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,5 +1,3 @@
-# src/streamlit_app.py
-
 import os
 import json
 import numpy as np
@@ -7,155 +5,142 @@ import streamlit as st
 from PIL import Image
 import tensorflow as tf
 
-# ==========================================================
+# --------------------------------------------------
 # Page config
-# ==========================================================
+# --------------------------------------------------
 st.set_page_config(
     page_title="ResNet50 Image Predictor",
     page_icon="🧠",
-    layout="centered",
+    layout="centered"
 )
 
 st.title("🧠 ResNet50 Image Predictor")
-st.write("Classifies mushroom images using a trained ResNet50 model.")
+st.write("Classifies mushroom images using a trained ResNet50 model (architecture in code + weights from src/).")
 
-# ==========================================================
-# Paths (HF Space: repo root is /app, your files are in /app/src)
-# ==========================================================
-MODEL_PATH = "src/new_best3_resnet50.keras"
-WEIGHTS_PATH = "src/resnet50_weights_noBN3.h5"
+# --------------------------------------------------
+# Paths (fixed, HF friendly)
+# --------------------------------------------------
+MODEL_WEIGHTS_PATH = "src/resnet50_weights_noBN3.h5"
 CLASS_NAMES_PATH = "src/class_names3.json"
 
 IMG_SIZE = (224, 224)
 
-# ==========================================================
+# --------------------------------------------------
 # Helpers
-# ==========================================================
-@st.cache_data
-def load_class_names(path: str):
+# --------------------------------------------------
+def _safe_listdir(path: str):
+    try:
+        return sorted(os.listdir(path))
+    except Exception as e:
+        return f"Could not list dir '{path}': {e}"
+
+def load_class_names(path: str) -> list:
     if not os.path.exists(path):
         raise FileNotFoundError(f"Missing class names file: {path}")
-    with open(path, "r") as f:
-        data = json.load(f)
-
-    # Accept either list ["a","b"] or dict {"0":"a","1":"b"}
-    if isinstance(data, dict):
-        # sort by numeric key if possible
-        try:
-            items = sorted(data.items(), key=lambda kv: int(kv[0]))
-            data = [v for _, v in items]
-        except Exception:
-            data = list(data.values())
-
-    if not isinstance(data, list) or len(data) == 0:
-        raise ValueError("class_names JSON must be a non-empty list (or dict).")
-
-    return data
-
+    with open(path, "r", encoding="utf-8") as f:
+        names = json.load(f)
+    if not isinstance(names, list) or len(names) == 0:
+        raise ValueError("class_names JSON must be a non-empty list.")
+    return names
+
+def build_resnet50_classifier(num_classes: int) -> tf.keras.Model:
+    base_model = tf.keras.applications.ResNet50(
+        weights="imagenet",
+        include_top=False,
+        input_shape=(IMG_SIZE[0], IMG_SIZE[1], 3),
+    )
+    base_model.trainable = False
+
+    model = tf.keras.Sequential([
+        base_model,
+        tf.keras.layers.GlobalAveragePooling2D(),
+        tf.keras.layers.Dense(256, activation="relu"),
+        tf.keras.layers.Dropout(0.5),
+        tf.keras.layers.Dense(num_classes, activation="softmax"),
+    ])
+
+    # Compile is not required for inference, but harmless:
+    model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
+    return model
 
 @st.cache_resource
-def load_model_and_weights(model_path: str, weights_path: str):
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"Missing model file: {model_path}")
+def load_model_and_assets(weights_path: str, class_names_path: str):
+    class_names = load_class_names(class_names_path)
+    model = build_resnet50_classifier(num_classes=len(class_names))
+
     if not os.path.exists(weights_path):
         raise FileNotFoundError(f"Missing weights file: {weights_path}")
 
-    # Load architecture/model (compile False is fine for inference)
-    model = tf.keras.models.load_model(model_path, compile=False)
-
-    # Load trained weights (must match the loaded model)
+    # Load weights into the exact same architecture
     model.load_weights(weights_path)
 
-    return model
-
+    return model, class_names
 
-def preprocess_for_resnet50(pil_img: Image.Image) -> np.ndarray:
+def preprocess_image(pil_img: Image.Image) -> np.ndarray:
     img = pil_img.convert("RGB").resize(IMG_SIZE)
     x = np.array(img, dtype=np.float32)
     x = np.expand_dims(x, axis=0)
     x = tf.keras.applications.resnet50.preprocess_input(x)
     return x
 
-
-# ==========================================================
-# Debug info (helpful on HuggingFace)
-# ==========================================================
+# --------------------------------------------------
+# Debug info (HF)
+# --------------------------------------------------
 with st.expander("🔍 Debug info (HuggingFace check)"):
-    try:
-        st.write("Working directory:", os.getcwd())
-        st.write("Files in ./ :", os.listdir("."))
-        if os.path.exists("src"):
-            st.write("Files in src/:", os.listdir("src"))
-        else:
-            st.warning("Folder 'src' not found.")
-        st.write("MODEL_PATH exists:", os.path.exists(MODEL_PATH))
-        st.write("WEIGHTS_PATH exists:", os.path.exists(WEIGHTS_PATH))
-        st.write("CLASS_NAMES exists:", os.path.exists(CLASS_NAMES_PATH))
-        st.write("TensorFlow version:", tf.__version__)
-    except Exception as e:
-        st.exception(e)
-
-# ==========================================================
-# Load artifacts
-# ==========================================================
+    st.write("Files in repo root:", _safe_listdir("."))
+    st.write("Files in src/:", _safe_listdir("src"))
+    st.write("Weights exists:", os.path.exists(MODEL_WEIGHTS_PATH), "->", MODEL_WEIGHTS_PATH)
+    st.write("Class names exists:", os.path.exists(CLASS_NAMES_PATH), "->", CLASS_NAMES_PATH)
+    st.write("TensorFlow version:", tf.__version__)
+
+# --------------------------------------------------
+# Load model + assets
+# --------------------------------------------------
 try:
-    class_names = load_class_names(CLASS_NAMES_PATH)
+    model, class_names = load_model_and_assets(MODEL_WEIGHTS_PATH, CLASS_NAMES_PATH)
+    st.success("✅ Model + weights loaded successfully!")
 except Exception as e:
-    st.error("❌ Could not load class names.")
+    st.error("❌ Model could not be loaded.")
     st.exception(e)
     st.stop()
 
+# --------------------------------------------------
 # Show available classes
+# --------------------------------------------------
 with st.expander("🧪 Available mushroom classes (you can test these)"):
+    st.write(f"Total classes: **{len(class_names)}**")
     for i, name in enumerate(class_names):
-        st.write(f"**{i}** → {name}")
-
-try:
-    model = load_model_and_weights(MODEL_PATH, WEIGHTS_PATH)
-    st.success("✅ Model + weights loaded successfully!")
-except Exception as e:
-    st.error("❌ Model could not be loaded.")
-    st.exception(e)
-    st.stop()
+        st.write(f"**{i}** — {name}")
 
-# ==========================================================
-# Upload + predict
-# ==========================================================
+# --------------------------------------------------
+# Image upload + prediction
+# --------------------------------------------------
 uploaded_file = st.file_uploader(
     "Upload a mushroom image",
-    type=["jpg", "jpeg", "png", "webp"],
+    type=["jpg", "jpeg", "png", "webp"]
 )
 
 if uploaded_file is None:
     st.info("👆 Please upload an image to start prediction.")
-    st.stop()
-
-# Show image
-img = Image.open(uploaded_file)
-st.image(img, caption="Uploaded image", use_container_width=True)
-
-# Predict
-x = preprocess_for_resnet50(img)
-preds = model.predict(x, verbose=0)[0]  # shape: (num_classes,)
-
-pred_idx = int(np.argmax(preds))
-pred_name = class_names[pred_idx] if pred_idx < len(class_names) else f"class_{pred_idx}"
-confidence = float(preds[pred_idx])
-
-st.subheader("✅ Prediction")
-st.write(f"**Predicted class index:** {pred_idx}")
-st.write(f"**Predicted mushroom:** 🍄 **{pred_name}**")
-st.write(f"**Confidence:** {confidence:.4f}")
-
-# Top-3
-st.subheader("📊 Top-3 predictions")
-top3_idx = np.argsort(preds)[::-1][:3]
-for i in top3_idx:
-    name = class_names[i] if i < len(class_names) else f"class_{i}"
-    st.write(f"**{name}** → {preds[i]:.4f}")
-
-# Full probs (optional)
-with st.expander("All class probabilities"):
-    for i, p in enumerate(preds):
-        name = class_names[i] if i < len(class_names) else f"class_{i}"
-        st.write(f"{i} → {name}: {p:.6f}")
+else:
+    img = Image.open(uploaded_file)
+    st.image(img, caption="Uploaded image", use_container_width=True)
+
+    x = preprocess_image(img)
+
+    preds = model.predict(x, verbose=0)[0]  # shape: (num_classes,)
+    pred_idx = int(np.argmax(preds))
+    pred_conf = float(preds[pred_idx])
+    pred_name = class_names[pred_idx] if pred_idx < len(class_names) else f"Class {pred_idx}"
+
+    st.subheader("✅ Prediction")
+    st.write(f"**Predicted class index:** {pred_idx}")
+    st.write(f"**Predicted class name:** {pred_name}")
+    st.write(f"**Confidence:** {pred_conf:.4f}")
+
+    st.subheader("🏁 Top-3 predictions")
+    top3_idx = np.argsort(preds)[::-1][:3]
+    for rank, idx in enumerate(top3_idx, start=1):
+        name = class_names[int(idx)]
+        prob = float(preds[int(idx)])
+        st.write(f"{rank}. **{name}** (class {int(idx)}) — **{prob*100:.2f}%**")