Update app.py

app.py

@@ -1,81 +1,110 @@
-import gradio as gr
-import pandas as pd
-import numpy as np
-import joblib
-
-# =========================
-# LOAD MODEL
-# =========================
-artifact = joblib.load("stacking_model.pkl")
-
-base_models   = artifact["base_models"]    # list of (name, model)
-meta_model    = artifact["meta_model"]
-FEATURE_NAMES = artifact["features"]
-
-N_FEATURES = len(FEATURE_NAMES)
-
-# =========================
-# PREDICTION FUNCTION
-# =========================
-def predict_csv(file):
-    df = pd.read_csv(file)
-
-    # Chuẩn hóa tên cột để tránh BOM/khoảng trắng
-    df.columns = df.columns.str.strip()
-    df.columns = df.columns.str.replace("\ufeff", "")
-
-    # Drop label columns if exist
-    df = df.drop(columns=["Label", "label", "class", "Class"], errors="ignore")
-
-    # Check missing features
-    missing_features = [f for f in FEATURE_NAMES if f not in df.columns]
-    if missing_features:
-        # Trả về DataFrame báo lỗi thay vì string
-        return pd.DataFrame({
-            "error": [f"Missing required features: {missing_features}"]
-        })
-
-    # Keep only required features & correct order
-    X = df[FEATURE_NAMES].astype(float)
-
-    # =========================
-    # LEVEL-1 (BASE MODELS)
-    # =========================
-    meta_inputs = []
-    for name, model in base_models:
-        probs = model.predict_proba(X)[:, 1]
-        meta_inputs.append(probs)
-
-    meta_X = np.column_stack(meta_inputs)
-
-    # =========================
-    # META MODEL
-    # =========================
-    final_probs = meta_model.predict_proba(meta_X)[:, 1]
-    final_preds = (final_probs > 0.5).astype(int)
-
-    # =========================
-    # BUILD OUTPUT
-    # =========================
-    result = df.copy()
-    result.insert(0, "row_id", range(1, len(df) + 1))
-    result["probability_malware"] = final_probs
-    result["prediction"] = final_preds
-    result["prediction_label"] = result["prediction"].map(
-        {1: "malware", 0: "benign"}
-    )
-
-    return result
-
-# =========================
-# GRADIO INTERFACE
-# =========================
-demo = gr.Interface(
-    fn=predict_csv,
-    inputs=gr.File(label="Upload CSV file"),
-    outputs=gr.Dataframe(label="Prediction Result"),
-    title="Stacking-based Malware Detection",
-    description="ExtraTrees + RandomForest + LightGBM + LogisticRegression → XGBoost"
-)
-
-demo.launch()
+import joblib
+import pandas as pd
+import gradio as gr
+import numpy as np
+import re
+
+# ======================
+# LOAD MODEL
+# ======================
+artifact = joblib.load("stacking_model.pkl")
+
+base_models = artifact["base_models"]      # list of (name, model)
+meta_model = artifact["meta_model"]
+feature_names = artifact["features"]
+
+# ======================
+# CLEAN FUNCTION (same as training)
+# ======================
+def clean_numeric(val):
+    if pd.isna(val):
+        return None
+
+    val = str(val).strip()
+    val = re.sub(r'\s+', '', val)
+
+    # scientific notation
+    if re.match(r'^-?\d+(\.\d+)?[eE][+-]?\d+$', val):
+        return float(val)
+
+    # remove thousand separators
+    if val.count('.') > 1:
+        val = val.replace('.', '')
+
+    # comma decimal -> dot
+    if ',' in val and '.' not in val:
+        val = val.replace(',', '.')
+
+    try:
+        return float(val)
+    except ValueError:
+        return None
+
+
+# ======================
+# PREDICTION FUNCTION
+# ======================
+def predict_malware_csv(file):
+    df = pd.read_csv(file.name)
+
+    # Check missing features
+    missing = set(feature_names) - set(df.columns)
+    if missing:
+        return f"❌ Missing features: {list(missing)}", None
+
+    # Keep only needed features
+    X = df[feature_names].copy()
+
+    # 🔥 CLEAN NUMERIC FEATURES
+    for col in feature_names:
+        X[col] = X[col].apply(clean_numeric)
+
+    # Optional: fill NaN if needed
+    # X = X.fillna(0)
+
+    # Level-1 predictions
+    meta_inputs = []
+    for name, model in base_models:
+        prob = model.predict_proba(X)[:, 1]
+        meta_inputs.append(prob)
+
+    meta_X = np.column_stack(meta_inputs)
+
+    # Meta prediction
+    preds = meta_model.predict(meta_X)
+    probs = meta_model.predict_proba(meta_X)[:, 1]
+
+    # Append results
+    result_df = df.copy()
+    result_df["Prediction"] = np.where(preds == 1, "Malware", "Benign")
+    result_df["Malware_Probability"] = probs
+
+    return "✅ Prediction completed", result_df
+
+
+# ======================
+# UI
+# ======================
+inputs = gr.File(
+    label="Upload CSV file (features only)",
+    file_types=[".csv"]
+)
+
+outputs = [
+    gr.Textbox(label="Status"),
+    gr.Dataframe(label="Prediction Results")
+]
+
+app = gr.Interface(
+    fn=predict_malware_csv,
+    inputs=inputs,
+    outputs=outputs,
+    title="Stacking-based Malware Detection",
+    description=(
+        "Upload a CSV file.\n\n"
+        f"Required features: {', '.join(feature_names)}"
+    )
+)
+
+if __name__ == "__main__":
+    app.launch()