Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,20 +1,22 @@
-# src/streamlit_app.py
-import os
-import json
-import pickle
-import numpy as np
 import streamlit as st
+import numpy as np
 from PIL import Image
-
 import tensorflow as tf
-from tensorflow.keras import Sequential
-from tensorflow.keras.layers import GlobalAveragePooling2D, BatchNormalization, Dense, Dropout
+import json
+import os
+
 from tensorflow.keras.applications import ResNet50
-from tensorflow.keras.applications.resnet50 import preprocess_input
+from tensorflow.keras.layers import (
+    Dense,
+    Dropout,
+    GlobalAveragePooling2D,
+    BatchNormalization
+)
+from tensorflow.keras.models import Sequential
 
-# -----------------------------
+# --------------------------------------------------
 # Page config
-# -----------------------------
+# --------------------------------------------------
 st.set_page_config(
     page_title="ResNet50 Image Predictor",
     page_icon="🧠",
@@ -24,95 +26,37 @@ st.set_page_config(
 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"
+# --------------------------------------------------
+# Paths (FIXED NAMES)
+# --------------------------------------------------
+WEIGHTS_PATH = "src/resnet50_weights1.h5"
+CLASS_NAMES_PATH = "src/class_names.json"
 
-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():
+# --------------------------------------------------
+# Debug info (VERY IMPORTANT on HuggingFace)
+# --------------------------------------------------
+with st.expander("🔍 Debug info (HuggingFace check)"):
+    st.write("Files in src/:")
     try:
-        return os.listdir("src")
+        st.write(os.listdir("src"))
+        st.write("Weights exist:", os.path.exists(WEIGHTS_PATH))
+        st.write("Classes exist:", os.path.exists(CLASS_NAMES_PATH))
     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!
-    """
+        st.error(e)
+
+# --------------------------------------------------
+# Load class names
+# --------------------------------------------------
+with open(CLASS_NAMES_PATH, "r") as f:
+    class_names = json.load(f)
+
+NUM_CLASSES = len(class_names)
+
+# --------------------------------------------------
+# Build model + load weights (HF SAFE)
+# --------------------------------------------------
+@st.cache_resource
+def load_model():
     base_model = ResNet50(
         weights="imagenet",
         include_top=False,
@@ -126,120 +70,60 @@ def build_resnet50_head(num_classes: int):
         BatchNormalization(),
         Dense(256, activation="relu"),
         Dropout(0.5),
-        Dense(num_classes, activation="softmax")
+        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))
+    model.load_weights(WEIGHTS_PATH)
+    return model
 
-    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
-# -----------------------------
+# --------------------------------------------------
+# Load model safely
+# --------------------------------------------------
 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}")
+    model = load_model()
+    st.success("✅ Model loaded successfully!")
 except Exception as e:
     st.error("❌ Model could not be loaded.")
     st.exception(e)
     st.stop()
 
-# -----------------------------
-# Upload & Predict
-# -----------------------------
+# --------------------------------------------------
+# Image upload
+# --------------------------------------------------
 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}")
+if uploaded_file is not None:
+    img = Image.open(uploaded_file).convert("RGB")
+    st.image(img, caption="Uploaded image", use_container_width=True)
+
+    # --------------------------------------------------
+    # Preprocessing (ResNet50)
+    # --------------------------------------------------
+    img = img.resize((224, 224))
+    x = np.array(img, dtype=np.float32)
+    x = np.expand_dims(x, axis=0)
+    x = tf.keras.applications.resnet50.preprocess_input(x)
+
+    # --------------------------------------------------
+    # Prediction
+    # --------------------------------------------------
+    preds = model.predict(x, verbose=0)
+    pred_idx = int(np.argmax(preds))
+    confidence = float(np.max(preds))
+
+    # --------------------------------------------------
+    # Output
+    # --------------------------------------------------
+    st.subheader("✅ Prediction")
+    st.write(f"**Class:** {class_names[pred_idx]}")
+    st.write(f"**Confidence:** {confidence:.4f}")
+
+    st.subheader("📊 Class probabilities")
+    for i, p in enumerate(preds[0]):
+        st.write(f"{class_names[i]}: {p:.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}")
+    st.info("👆 Please upload an image to start prediction.")