Update app.py

app.py

@@ -4,36 +4,16 @@ import numpy as np
 import tensorflow as tf
 from sklearn.preprocessing import MinMaxScaler
 import plotly.express as px
+import os
 from groq import Groq
-import io
 
 # Initialize session state
 if 'model' not in st.session_state:
     st.session_state.model = None
 if 'threshold' not in st.session_state:
     st.session_state.threshold = None
-if 'data' not in st.session_state:
-    st.session_state.data = None
 
-# Generate synthetic data
-def generate_sample_data():
-    time_steps = 500
-    base = np.arange(time_steps)
-    data = pd.DataFrame({
-        'timestamp': pd.date_range(start='2024-01-01', periods=time_steps, freq='H'),
-        'device_count': np.random.poisson(50, time_steps),
-        'connection_attempts': np.random.poisson(30, time_steps),
-        'packet_loss': np.random.uniform(0.1, 2.0, time_steps),
-        'latency': np.random.uniform(10, 100, time_steps)
-    })
-    
-    # Add anomalies
-    anomaly_indices = np.random.choice(time_steps, 10, replace=False)
-    data.loc[anomaly_indices, 'connection_attempts'] *= 10
-    data.loc[anomaly_indices, 'latency'] *= 5
-    return data
-
-# Autoencoder model
+# Autoencoder model definition
 def build_autoencoder(input_dim):
     model = tf.keras.Sequential([
         tf.keras.layers.Dense(64, activation='relu', input_shape=(input_dim,)),
@@ -51,96 +31,74 @@ groq_api_key = st.sidebar.text_input("Groq API Key (optional)", type="password")
 
 # Main interface
 st.title("🛰️ AI Network Anomaly Detector")
-st.write("Upload network data (CSV) or use synthetic data")
+st.write("Upload your network data (CSV) to detect anomalies")
 
-# Data handling
+# File uploader
 uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
-if uploaded_file:
-    try:
-        st.session_state.data = pd.read_csv(uploaded_file)
-    except Exception as e:
-        st.error(f"Error reading file: {str(e)}")
+
+# Load or generate sample data
+if uploaded_file is not None:
+    data = pd.read_csv(uploaded_file)
 else:
-    st.session_state.data = generate_sample_data()
-    st.info("Using synthetic data. Upload a CSV file to analyze your own data.")
+    st.info("Using sample data. Upload a file to use your own dataset.")
+    data = pd.read_csv("sample_wifi_data.csv")  # You should provide this sample file
 
 # Preprocessing
-if st.session_state.data is not None:
-    features = ['device_count', 'connection_attempts', 'packet_loss', 'latency']
-    scaler = MinMaxScaler()
-    data_scaled = scaler.fit_transform(st.session_state.data[features])
+features = ['device_count', 'connection_attempts', 'packet_loss', 'latency']
+scaler = MinMaxScaler()
+data_scaled = scaler.fit_transform(data[features])
 
-# Model training
-if st.session_state.model is None or fine_tune:
+# Model training/fine-tuning
+if fine_tune or st.session_state.model is None:
     with st.spinner("Training model..."):
-        try:
-            autoencoder = build_autoencoder(data_scaled.shape[1])
-            autoencoder.fit(data_scaled, data_scaled,
-                          epochs=100,
-                          batch_size=32,
-                          verbose=0,
-                          validation_split=0.2)
-            st.session_state.model = autoencoder
-            
-            # Calculate threshold
-            reconstructions = autoencoder.predict(data_scaled)
-            mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1)
-            st.session_state.threshold = np.percentile(mse, 95)
-            st.success("Model trained successfully!")
-        except Exception as e:
-            st.error(f"Training error: {str(e)}")
+        autoencoder = build_autoencoder(data_scaled.shape[1])
+        autoencoder.fit(data_scaled, data_scaled, 
+                       epochs=100, 
+                       batch_size=32,
+                       verbose=0,
+                       validation_split=0.1)
+        st.session_state.model = autoencoder
+        
+        # Calculate threshold
+        reconstructions = autoencoder.predict(data_scaled)
+        mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1)
+        st.session_state.threshold = np.percentile(mse, 95)
 
 # Anomaly detection
 if st.session_state.model and st.button("Detect Anomalies"):
-    try:
-        reconstructions = st.session_state.model.predict(data_scaled)
-        mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1)
-        anomalies = mse > st.session_state.threshold
-        
-        # Visualization
-        fig = px.line(st.session_state.data, x='timestamp', y='connection_attempts',
-                     title='Network Traffic with Anomalies')
-        fig.add_scatter(x=st.session_state.data[anomalies]['timestamp'],
-                       y=st.session_state.data[anomalies]['connection_attempts'],
-                       mode='markers', name='Anomalies',
-                       marker=dict(color='red', size=8))
-        st.plotly_chart(fig)
-
-        # Groq API integration
-        if groq_api_key:
-            try:
-                client = Groq(api_key=groq_api_key)
-                response = client.chat.completions.create(
-                    model="llama3-70b-8192",
-                    messages=[{
-                        "role": "user",
-                        "content": f"""Network anomaly report:
-                        - Total data points: {len(st.session_state.data)}
-                        - Anomalies detected: {sum(anomalies)}
-                        - Max connection attempts: {st.session_state.data['connection_attempts'].max()}
-                        - Average latency: {st.session_state.data['latency'].mean():.2f}ms
-                        Provide a technical analysis and recommendations in bullet points."""
-                    }]
-                )
-                st.subheader("AI Analysis")
-                st.write(response.choices[0].message.content)
-            except Exception as e:
-                st.error(f"Groq API Error: {str(e)}")
-        else:
-            st.warning("Groq API key not provided - using basic detection")
-
-        # Download report
-        report = st.session_state.data[anomalies]
-        csv = report.to_csv(index=False).encode('utf-8')
-        st.download_button(
-            label="Download Anomaly Report",
-            data=csv,
-            file_name='anomaly_report.csv',
-            mime='text/csv'
-        )
-    except Exception as e:
-        st.error(f"Detection error: {str(e)}")
+    reconstructions = st.session_state.model.predict(data_scaled)
+    mse = np.mean(np.power(data_scaled - reconstructions, 2), axis=1)
+    anomalies = mse > st.session_state.threshold
+    
+    # Visualization
+    fig = px.line(data, x=data.index, y='connection_attempts', 
+                 title='Network Traffic with Anomalies')
+    fig.add_scatter(x=data[anomalies].index, y=data[anomalies]['connection_attempts'], 
+                   mode='markers', name='Anomalies')
+    st.plotly_chart(fig)
+    
+    # Generate alert with Groq/Llama3
+    if groq_api_key:
+        try:
+            client = Groq(api_key=groq_api_key)
+            response = client.chat.completions.create(
+                model="llama3-70b-8192",
+                messages=[{
+                    "role": "user",
+                    "content": f"Generate a network security alert for {sum(anomalies)} anomalies detected. Max connection attempts: {data['connection_attempts'].max()}"
+                }]
+            )
+            st.warning(response.choices[0].message.content)
+        except Exception as e:
+            st.error(f"Groq API Error: {str(e)}")
+    else:
+        st.warning(f"Detected {sum(anomalies)} anomalies! Consider adding Groq API key for detailed analysis.")
 
-# Display raw data
-if st.checkbox("Show raw data"):
-    st.write(st.session_state.data)
+# Download button for results
+if st.session_state.threshold:
+    st.download_button(
+        label="Download Anomaly Report",
+        data=data[anomalies].to_csv().encode('utf-8'),
+        file_name='anomalies_report.csv',
+        mime='text/csv'
+    )