Update app.py

app.py

@@ -1,21 +1,54 @@
-import gradio as gr
-from utils import load_all_models, predict_image
-
-# Load all models and class mappings once at startup
-models, offsets, idx_to_class = load_all_models()
-
-# Inference function wrapper
-def classify(image):
-    return predict_image(image, models, offsets, idx_to_class)
-
-# Launch Gradio interface
-demo = gr.Interface(
-    fn=classify,
-    inputs=gr.Image(type="pil", label="Upload Weed Image"),
-    outputs=gr.Label(num_top_classes=3, label="Top Predicted Classes"),
-    title="🌿 Weed Species Classifier",
-    description="Upload a weed image to classify it into one of 25 species using confidence-based routing across 3 Swin Transformer models.",
-    live=True  # ✅ Shows a spinner while running
-)
-
-demo.launch()
+import streamlit as st
+from utils import load_model, preprocess_image, predict_final_class
+
+st.set_page_config(page_title="Weed Classifier", layout="centered")
+
+st.title("🌿 Weed Species Classifier")
+st.write("Upload an image of a weed, and the model will classify it.")
+
+# Define the models and class mappings
+model_defs = [
+    {
+        "name": "Model 1",
+        "path": "MMIM_best1.pth",
+        "class_names": ["class10", "class11", "class12", "class13"]
+    },
+    {
+        "name": "Model 2",
+        "path": "MMIM_best2.pth",
+        "class_names": ["class14", "class15", "class16", "class17", "class18", "class19"]
+    },
+    {
+        "name": "Model 3",
+        "path": "MMIM_best3.pth",
+        "class_names": ["class20", "class21", "class22", "class23", "class24", "class25"]
+    }
+]
+
+# Load models once (on app startup)
+@st.cache_resource
+def load_all_models():
+    for model_def in model_defs:
+        model_def["model"] = load_model(model_def["path"], len(model_def["class_names"]))
+    return model_defs
+
+model_defs = load_all_models()
+
+# Upload section
+uploaded_image = st.file_uploader("📤 Upload Weed Image", type=["jpg", "jpeg", "png"])
+
+# Prediction
+if uploaded_image:
+    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
+
+    # Preprocess and predict
+    image_tensor = preprocess_image(uploaded_image)
+    predicted_class = predict_final_class(image_tensor, model_defs)
+
+    # Display result
+    st.markdown("## 🔍 Predicted Class")
+    st.success(f"**{predicted_class}**")
+
+    st.markdown("---")
+    if st.button("Clear"):
+        st.experimental_rerun()