hari6677 commited on
Commit
a0f7d92
·
verified ·
1 Parent(s): 5425f37

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +135 -0
app.py ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app.py
2
+
3
+ import os
4
+ import numpy as np
5
+ import pandas as pd
6
+ import tensorflow as tf
7
+ from tensorflow.keras.models import load_model
8
+ import pickle
9
+ from flask import Flask, request, jsonify, render_template
10
+
11
+ # --- Configuration ---
12
+ MODEL_PATH = 'cicids2017_mlp_model.keras'
13
+ SCALER_PATH = 'cicids2017_scaler.pkl'
14
+ EXPECTED_FEATURES = 49 # Based on your training data shape
15
+ app = Flask(__name__)
16
+
17
+ # --- Global Assets ---
18
+ loaded_model = None
19
+ loaded_scaler = None
20
+
21
+ def load_assets():
22
+ """Load the model and scaler only once when the app starts."""
23
+ global loaded_model, loaded_scaler
24
+
25
+ # Check if files exist
26
+ if not os.path.exists(MODEL_PATH) or not os.path.exists(SCALER_PATH):
27
+ print("🚨 ERROR: Model or scaler files not found. Ensure they are in the same directory.")
28
+ return False
29
+
30
+ try:
31
+ # Load the Keras model
32
+ loaded_model = load_model(MODEL_PATH)
33
+
34
+ # Load the StandardScaler object
35
+ with open(SCALER_PATH, 'rb') as file:
36
+ loaded_scaler = pickle.load(file)
37
+
38
+ print(f"✅ Assets loaded successfully. Model ready for {EXPECTED_FEATURES} features.")
39
+ return True
40
+
41
+ except Exception as e:
42
+ print(f"🚨 FATAL ERROR loading assets: {e}")
43
+ return False
44
+
45
+ # Load assets when the application starts
46
+ load_assets()
47
+
48
+
49
+ # --- HTML Template for Simple Interface ---
50
+ HTML_TEMPLATE = """
51
+ <!doctype html>
52
+ <title>NIDS Prediction</title>
53
+ <style>
54
+ body {{ font-family: sans-serif; max-width: 800px; margin: auto; padding: 20px; }}
55
+ textarea {{ width: 100%; min-height: 150px; padding: 10px; box-sizing: border-box; }}
56
+ h1 {{ color: #007bff; }}
57
+ .result {{ padding: 15px; border-radius: 5px; margin-top: 20px; }}
58
+ .attack {{ background-color: #f8d7da; border: 1px solid #f5c6cb; color: #721c24; }}
59
+ .benign {{ background-color: #d4edda; border: 1px solid #c3e6cb; color: #155724; }}
60
+ .error {{ background-color: #ffeeba; border: 1px solid #ffc720; color: #664d03; }}
61
+ </style>
62
+ <h1>Network Intrusion Detection System (NIDS)</h1>
63
+ <p>Paste {num_features} comma-separated network flow features below for prediction.</p>
64
+
65
+ <form method="POST" action="/">
66
+ <textarea name="features" placeholder="e.g., 0.1, 10.5, 0.003, ...">{example_data}</textarea><br><br>
67
+ <input type="submit" value="Predict Traffic Type" style="padding: 10px 20px; background-color: #007bff; color: white; border: none; cursor: pointer;">
68
+ </form>
69
+
70
+ {result_html}
71
+ """
72
+
73
+
74
+ # --- Flask Routes ---
75
+
76
+ @app.route('/', methods=['GET', 'POST'])
77
+ def home():
78
+ """Handles both the form display (GET) and prediction submission (POST)."""
79
+ result_html = ""
80
+ example_data = "" # You can load a real example here if you have a local test file
81
+
82
+ if request.method == 'POST':
83
+ try:
84
+ # 1. Get raw input data
85
+ features_str = request.form.get('features')
86
+ features_list = [float(x.strip()) for x in features_str.split(',') if x.strip()]
87
+ example_data = features_str
88
+
89
+ if len(features_list) != EXPECTED_FEATURES:
90
+ raise ValueError(f"Input must contain exactly {EXPECTED_FEATURES} features, but received {len(features_list)}.")
91
+
92
+ # 2. Scale the input
93
+ input_array = np.array(features_list).reshape(1, -1)
94
+ input_scaled = loaded_scaler.transform(input_array)
95
+
96
+ # 3. Predict
97
+ prediction_proba = loaded_model.predict(input_scaled, verbose=0)[0][0]
98
+ prediction_class = 1 if prediction_proba > 0.5 else 0
99
+
100
+ # 4. Format Result
101
+ class_label = "🚨 ATTACK" if prediction_class == 1 else "✅ BENIGN"
102
+ css_class = "attack" if prediction_class == 1 else "benign"
103
+
104
+ result_html = f"""
105
+ <div class="result {css_class}">
106
+ <h3>Prediction: {class_label}</h3>
107
+ <p>Probability of Attack (1): <b>{prediction_proba:.5f}</b></p>
108
+ </div>
109
+ """
110
+
111
+ except ValueError as e:
112
+ result_html = f"""
113
+ <div class="result error">
114
+ <h3>Input Error</h3>
115
+ <p>{e}</p>
116
+ </div>
117
+ """
118
+ except Exception as e:
119
+ result_html = f"""
120
+ <div class="result error">
121
+ <h3>Prediction Error</h3>
122
+ <p>An unexpected error occurred: {e}</p>
123
+ </div>
124
+ """
125
+
126
+ return render_template_string(
127
+ HTML_TEMPLATE,
128
+ num_features=EXPECTED_FEATURES,
129
+ example_data=example_data,
130
+ result_html=result_html
131
+ )
132
+
133
+ if __name__ == '__main__':
134
+ # Use 0.0.0.0 for compatibility with Docker/server environments
135
+ app.run(host='0.0.0.0', port=5000)