|
|
from flask import Flask, render_template, request, jsonify
|
|
|
import torch
|
|
|
import pandas as pd
|
|
|
import numpy as np
|
|
|
import pickle
|
|
|
from torch_geometric.data import Data
|
|
|
import config
|
|
|
from model import ADSB_GAT
|
|
|
|
|
|
app = Flask(__name__)
|
|
|
|
|
|
# 全局变量
|
|
|
model = None
|
|
|
scaler = None
|
|
|
le_icao = None
|
|
|
FEATURE_COLS = None
|
|
|
|
|
|
def load_model_once():
|
|
|
global model, scaler, le_icao, FEATURE_COLS
|
|
|
print("正在加载模型和预处理参数...")
|
|
|
|
|
|
# 1. 加载训练好的GAT模型
|
|
|
model = ADSB_GAT(in_channels=8)
|
|
|
model.load_state_dict(torch.load(config.MODEL_SAVE_PATH, map_location=config.DEVICE))
|
|
|
model.to(config.DEVICE)
|
|
|
model.eval()
|
|
|
|
|
|
# 2. 加载预处理参数（和训练时完全一致）
|
|
|
with open("scaler.pkl", "rb") as f:
|
|
|
scaler = pickle.load(f)
|
|
|
with open("le_icao.pkl", "rb") as f:
|
|
|
le_icao = pickle.load(f)
|
|
|
with open("feature_cols.pkl", "rb") as f:
|
|
|
FEATURE_COLS = pickle.load(f)
|
|
|
|
|
|
print("✅ 加载完成！服务已就绪")
|
|
|
print(f" 模型输入特征：{FEATURE_COLS}")
|
|
|
|
|
|
@app.route('/')
|
|
|
def index():
|
|
|
return render_template('index.html')
|
|
|
|
|
|
@app.route('/predict', methods=['POST'])
|
|
|
def predict():
|
|
|
try:
|
|
|
# 1. 接收上传的文件
|
|
|
file = request.files['file']
|
|
|
if not file:
|
|
|
return jsonify({"error": "请先上传CSV文件"}), 400
|
|
|
|
|
|
# 2. 读取全量数据
|
|
|
df = pd.read_csv(file)
|
|
|
total_original = len(df)
|
|
|
|
|
|
# ====================== 【新增】只取前1万条数据 ======================
|
|
|
MAX_LINES = 10000
|
|
|
if total_original > MAX_LINES:
|
|
|
df = df.iloc[:MAX_LINES] # 只保留前10000行
|
|
|
# =======================================================================
|
|
|
|
|
|
df = df.dropna(axis=0)
|
|
|
total_valid = len(df)
|
|
|
|
|
|
if total_valid == 0:
|
|
|
return jsonify({"error": "文件中没有有效数据"}), 400
|
|
|
|
|
|
# 3. 预处理（和训练时完全一致的步骤）
|
|
|
# 3.1 编码ICAO24
|
|
|
if "icao24" in df.columns:
|
|
|
df["icao24"] = df["icao24"].astype(str)
|
|
|
known_icao = set(le_icao.classes_)
|
|
|
df["icao24_enc"] = df["icao24"].apply(lambda x: le_icao.transform([x])[0] if x in known_icao else 0)
|
|
|
else:
|
|
|
df["icao24_enc"] = 0
|
|
|
|
|
|
# 3.2 计算加速度衍生特征
|
|
|
df = df.sort_values(["time"]).reset_index(drop=True)
|
|
|
df['acceleration'] = df['velocity'].diff().fillna(0)
|
|
|
|
|
|
# 3.3 【核心修复】用完整8个特征做归一化，和训练时完全匹配
|
|
|
df[FEATURE_COLS] = scaler.transform(df[FEATURE_COLS])
|
|
|
|
|
|
# 4. 按时间窗口切片构图
|
|
|
df["time_window"] = (df["time"] // config.TIME_WINDOW).astype(int)
|
|
|
window_groups = df.groupby("time_window")
|
|
|
|
|
|
all_preds = []
|
|
|
all_labels = []
|
|
|
|
|
|
for window_id, window_df in window_groups:
|
|
|
n = len(window_df)
|
|
|
# 节点太少直接判定为正常
|
|
|
if n < 3:
|
|
|
all_preds.extend([0]*n)
|
|
|
if "label" in window_df.columns:
|
|
|
all_labels.extend(window_df["label"].tolist())
|
|
|
continue
|
|
|
|
|
|
# 构建节点特征
|
|
|
x = torch.tensor(window_df[FEATURE_COLS].values, dtype=torch.float).to(config.DEVICE)
|
|
|
|
|
|
# 构建边（和训练时一致的规则）
|
|
|
edge_index = []
|
|
|
for i in range(n):
|
|
|
for j in range(i+1, min(i+6, n)):
|
|
|
edge_index.append([i, j])
|
|
|
edge_index.append([j, i])
|
|
|
edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous().to(config.DEVICE)
|
|
|
|
|
|
# 模型推理
|
|
|
with torch.no_grad():
|
|
|
out, _ = model(x, edge_index)
|
|
|
preds = out.argmax(dim=1).cpu().numpy().tolist()
|
|
|
|
|
|
all_preds.extend(preds)
|
|
|
# 如果有真实标签，同步保存
|
|
|
if "label" in window_df.columns:
|
|
|
all_labels.extend(window_df["label"].tolist())
|
|
|
|
|
|
# 5. 统计结果
|
|
|
total = len(all_preds)
|
|
|
attack_count = int(np.sum(np.array(all_preds) == 1))
|
|
|
normal_count = total - attack_count
|
|
|
attack_rate = round(attack_count / total * 100, 2)
|
|
|
|
|
|
# 有真实标签的话，计算准确率
|
|
|
accuracy = None
|
|
|
if all_labels and len(all_labels) == total:
|
|
|
correct = np.sum(np.array(all_preds) == np.array(all_labels))
|
|
|
accuracy = round(correct / total * 100, 2)
|
|
|
|
|
|
return jsonify({
|
|
|
"status": "success",
|
|
|
"total_original": total_original,
|
|
|
"total_valid": total,
|
|
|
"normal": normal_count,
|
|
|
"attack": attack_count,
|
|
|
"attack_rate": attack_rate,
|
|
|
"accuracy": accuracy
|
|
|
})
|
|
|
|
|
|
except Exception as e:
|
|
|
return jsonify({"error": f"处理失败：{str(e)}"}), 500
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
load_model_once()
|
|
|
print("服务启动成功！访问 http://localhost:5000")
|
|
|
app.run(debug=True, host='0.0.0.0', port=5000) |
