Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,20 +8,19 @@ import tensorflow as tf
 from tensorflow.keras.models import load_model
 from tensorflow.keras.optimizers.legacy import SGD
 
-# ==== File Paths (update if needed) ====
-
-MODEL_FILE =  ("src/model1.h5")
-WEIGHTS_FILE = ("src/weights.h5")
-SCALER_FILE =  ("src/standard_scaler1.pkl")
-LABEL_ENCODER_FILE = ("src/label_encoder1.pkl")
+# ==== File Paths ====
+MODEL_FILE = "src/model1.h5"
+WEIGHTS_FILE = "src/weights.h5"
+SCALER_FILE = "src/standard_scaler1.pkl"
+LABEL_ENCODER_FILE = "src/label_encoder1.pkl"
 
 ENCODER_PATHS = {
-    "proto":  ("src/categorical_label_encoder_proto1.pkl"),
-    "conn_state":  ("src/categorical_label_encoder_conn_state1.pkl"),
-    "history":  ("src/categorical_label_encoder_history1.pkl")
+    "proto": "src/categorical_label_encoder_proto1.pkl",
+    "conn_state": "src/categorical_label_encoder_conn_state1.pkl",
+    "history": "src/categorical_label_encoder_history1.pkl"
 }
 
-SAMPLE_FILE = "src/sample_input.csv" 
+SAMPLE_FILE = "src/sample_input.csv"
 
 # ==== Class for Malware Prediction ====
 class MalwareClassifier:
@@ -101,15 +100,10 @@ class MalwareClassifier:
 
 # ==== Streamlit UI ====
 def main():
-    import os
     st.set_page_config(page_title="Malware Detection System", page_icon="🛡️")
     st.title("🛡️ Malware Detection System")
     st.markdown("Upload a CSV file with network traffic logs to detect malware.")
 
-    # 🛠️ DEBUG INFO
-   
-    # 🔽 Sample CSV download
-    SAMPLE_FILE = "src/sample_input.csv"
     st.markdown("📄 Need a sample file to test? Download below:")
 
     if os.path.exists(SAMPLE_FILE):
@@ -117,11 +111,11 @@ def main():
             st.download_button(
                 label="📥 Download Sample CSV",
                 data=f,
-                file_name="src/sample_input.csv",
+                file_name="sample_input.csv",  # ✅ Do NOT include `src/` in file_name
                 mime="text/csv"
             )
     else:
-        st.warning("⚠️ Sample CSV not found. Please place 'sample_input.csv' in the app directory.")
+        st.warning("⚠️ 'sample_input.csv' not found in the 'src' directory.")
 
     try:
         classifier = MalwareClassifier()
@@ -129,12 +123,12 @@ def main():
         st.error(f"❌ Model loading failed: {e}")
         return
 
-    uploaded_file = st.file_uploader("📂 Drag & drop or select a CSV file", type=["csv"])
+    uploaded_file = st.file_uploader("📂 Upload your network CSV", type=["csv"])
 
     if uploaded_file is not None:
         try:
-            df = pd.read_csv(uploaded_file, delimiter=',')
-            df.columns = df.columns.str.strip()  # ✅ Normalize column names
+            df = pd.read_csv(uploaded_file)
+            df.columns = df.columns.str.strip()
 
             required_prediction_columns = [
                 'id.orig_h', 'id.orig_p', 'id.resp_h', 'id.resp_p',
@@ -142,21 +136,16 @@ def main():
                 'orig_pkts', 'orig_ip_bytes', 'resp_pkts', 'resp_ip_bytes'
             ]
 
-            # ✅ Check for missing columns
             missing_columns = set(required_prediction_columns) - set(df.columns)
             if missing_columns:
                 st.error(f"❌ Missing required columns for prediction: {missing_columns}")
-                st.write("📋 Detected columns:", df.columns.tolist())  # Helpful debug info
+                st.write("📋 Detected columns:", df.columns.tolist())
                 return
 
-            # ✅ Extract necessary columns
             prediction_input = df[required_prediction_columns].copy()
-
-            # ✅ Predict
             predictions = classifier.predict(prediction_input)
             st.success("✅ Prediction complete!")
 
-            # ✅ Display results
             for i, result in enumerate(predictions):
                 st.subheader(f"Prediction {i + 1}")
                 st.write(f"**Predicted Label:** {result['result']}")