Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,245 @@
-import altair as alt
+# src/streamlit_app.py
+import os
+import json
+import pickle
 import numpy as np
-import pandas as pd
 import streamlit as st
+from PIL import Image
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+import tensorflow as tf
+from tensorflow.keras import Sequential
+from tensorflow.keras.layers import GlobalAveragePooling2D, BatchNormalization, Dense, Dropout
+from tensorflow.keras.applications import ResNet50
+from tensorflow.keras.applications.resnet50 import preprocess_input
+
+# -----------------------------
+# Page config
+# -----------------------------
+st.set_page_config(
+    page_title="ResNet50 Image Predictor",
+    page_icon="🧠",
+    layout="centered"
+)
+
+st.title("🧠 ResNet50 Image Predictor")
+st.write("Classifies mushroom images using a trained ResNet50 model.")
+
+# -----------------------------
+# Fixed paths (from src/)
+# -----------------------------
+MODEL_PATH = "src/new_best2_resnet50.keras"
+WEIGHTS_PATH = "src/resnet50_weights2.h5"
+CLASS_NAMES_PATH = "src/class_names2"
+
+IMG_SIZE = (224, 224)
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def file_exists(path: str) -> bool:
+    try:
+        return os.path.exists(path)
+    except Exception:
+        return False
+
+def list_src_files():
+    try:
+        return os.listdir("src")
+    except Exception as e:
+        return [f"Could not list src/: {e}"]
+
+def load_class_names(path_base: str):
+    """
+    Loads class names from:
+    - class_names2 (no extension) as: pickle / json / txt
+    - class_names2.pkl / .json / .txt / .npy
+    If nothing works, returns None.
+    """
+    candidates = [
+        path_base,
+        path_base + ".pkl",
+        path_base + ".pickle",
+        path_base + ".json",
+        path_base + ".txt",
+        path_base + ".npy",
+    ]
+
+    for p in candidates:
+        if not file_exists(p):
+            continue
+
+        # numpy
+        if p.endswith(".npy"):
+            try:
+                arr = np.load(p, allow_pickle=True)
+                names = arr.tolist()
+                if isinstance(names, (list, tuple)) and len(names) > 0:
+                    return list(names), p
+            except Exception:
+                pass
+
+        # json
+        if p.endswith(".json"):
+            try:
+                with open(p, "r", encoding="utf-8") as f:
+                    names = json.load(f)
+                if isinstance(names, (list, tuple)) and len(names) > 0:
+                    return list(names), p
+            except Exception:
+                pass
+
+        # txt
+        if p.endswith(".txt"):
+            try:
+                with open(p, "r", encoding="utf-8") as f:
+                    names = [line.strip() for line in f.readlines() if line.strip()]
+                if len(names) > 0:
+                    return names, p
+            except Exception:
+                pass
+
+        # try pickle for "no extension" or .pkl/.pickle
+        try:
+            with open(p, "rb") as f:
+                names = pickle.load(f)
+            if isinstance(names, (list, tuple)) and len(names) > 0:
+                return list(names), p
+        except Exception:
+            pass
+
+    return None, None
+
+def build_resnet50_head(num_classes: int):
+    """
+    Must match your training architecture exactly!
+    """
+    base_model = ResNet50(
+        weights="imagenet",
+        include_top=False,
+        input_shape=(224, 224, 3)
+    )
+    base_model.trainable = False
+
+    model = Sequential([
+        base_model,
+        GlobalAveragePooling2D(),
+        BatchNormalization(),
+        Dense(256, activation="relu"),
+        Dropout(0.5),
+        Dense(num_classes, activation="softmax")
+    ])
+    return model
+
+@st.cache_resource
+def load_trained_model():
+    """
+    1) Try loading full .keras model
+    2) If fails -> rebuild architecture + load weights
+    """
+    # Load class names (needed for weights fallback)
+    class_names, class_src = load_class_names(CLASS_NAMES_PATH)
+
+    # ---- Try full model first
+    if file_exists(MODEL_PATH):
+        try:
+            m = tf.keras.models.load_model(MODEL_PATH, compile=False)
+            return m, class_names, f"Loaded full model: {MODEL_PATH}", class_src
+        except Exception as e:
+            full_model_error = str(e)
+    else:
+        full_model_error = f"Full model not found: {MODEL_PATH}"
+
+    # ---- Fallback: build architecture + load weights
+    if class_names is None:
+        raise RuntimeError(
+            "Full model loading failed AND class names could not be loaded.\n"
+            f"Full model error: {full_model_error}\n"
+            "Please upload class_names2 as .pkl or .txt or .json into src/."
+        )
+
+    if not file_exists(WEIGHTS_PATH):
+        raise RuntimeError(
+            "Full model loading failed AND weights file not found.\n"
+            f"Full model error: {full_model_error}\n"
+            f"Weights not found: {WEIGHTS_PATH}"
+        )
+
+    model = build_resnet50_head(num_classes=len(class_names))
+
+    try:
+        model.load_weights(WEIGHTS_PATH)
+        return model, class_names, f"Loaded weights: {WEIGHTS_PATH} (fallback)", class_src
+    except Exception as e:
+        raise RuntimeError(
+            "Full model loading failed AND weights loading failed.\n"
+            f"Full model error: {full_model_error}\n"
+            f"Weights error: {e}"
+        )
+
+# -----------------------------
+# Debug panel
+# -----------------------------
+with st.expander("🔍 Debug info (HuggingFace check)"):
+    st.write("Files in src/:")
+    st.write(list_src_files())
+    st.write("MODEL_PATH exists:", file_exists(MODEL_PATH))
+    st.write("WEIGHTS_PATH exists:", file_exists(WEIGHTS_PATH))
+    st.write("CLASS_NAMES base exists:", file_exists(CLASS_NAMES_PATH))
+    st.write("TensorFlow:", tf.__version__)
+
+# -----------------------------
+# Load model
+# -----------------------------
+try:
+    model, class_names, load_msg, class_src = load_trained_model()
+    st.success(f"✅ Model loaded. {load_msg}")
+    if class_names is not None:
+        st.caption(f"Class names loaded from: {class_src}")
+except Exception as e:
+    st.error("❌ Model could not be loaded.")
+    st.exception(e)
+    st.stop()
+
+# -----------------------------
+# Upload & Predict
+# -----------------------------
+uploaded_file = st.file_uploader(
+    "Upload a mushroom image",
+    type=["jpg", "jpeg", "png", "webp"]
+)
+
+if uploaded_file is None:
+    st.info("👆 Please upload an image to start prediction.")
+    st.stop()
+
+img = Image.open(uploaded_file).convert("RGB")
+st.image(img, caption="Uploaded image", use_container_width=True)
+
+# preprocess
+img_resized = img.resize(IMG_SIZE)
+x = np.array(img_resized, dtype=np.float32)
+x = np.expand_dims(x, axis=0)
+x = preprocess_input(x)
+
+# predict
+preds = model.predict(x, verbose=0)[0]
+idx = int(np.argmax(preds))
+conf = float(preds[idx])
+
+# output
+st.subheader("✅ Prediction")
+
+if class_names is not None and idx < len(class_names):
+    st.write(f"**Predicted class:** {class_names[idx]}")
+else:
+    st.write(f"**Predicted class index:** {idx}")
+
+st.write(f"**Confidence:** {conf:.4f}")
+
+st.subheader("📊 Class probabilities")
+if class_names is not None and len(class_names) == len(preds):
+    for name, p in zip(class_names, preds):
+        st.write(f"{name}: {float(p):.4f}")
+else:
+    for i, p in enumerate(preds):
+        st.write(f"Class {i}: {float(p):.4f}")