Upload app.py

app.py

@@ -0,0 +1,104 @@
+import gradio as gr
+import pandas as pd
+import numpy as np
+import joblib
+import tensorflow as tf
+
+# =========================
+# LOAD MODEL & SCALER
+# =========================
+model = tf.keras.models.load_model("mlp_malware.keras")
+scaler = joblib.load("scaler.pkl")
+
+# =========================
+# 30 SELECTED FEATURES
+
+SELECTED_FEATURES = [
+    "filesize",
+    "E_file",
+    "E_text",
+    "E_data",
+    "AddressOfEntryPoint",
+    "NumberOfSections",
+    "SizeOfInitializedData",
+    "SizeOfImage",
+    "SizeOfOptionalHeader",
+    "SizeOfCode",
+    "DirectoryEntryImportSize",
+    "ImageBase",
+    "CheckSum",
+    "Magic",
+    "MinorLinkerVersion",
+    "MajorSubsystemVersion",
+    "e_lfanew",
+    "sus_sections",
+    "PointerToSymbolTable",
+    "SectionsLength",
+    "SizeOfStackReserve",
+    "MajorOperatingSystemVersion",
+    "non_sus_sections",
+    "Characteristics",
+    "NumberOfSymbols",
+    "BaseOfData",
+    "MajorImageVersion",
+    "FH_char5",
+    "FH_char8",
+    "OH_DLLchar5"
+]
+
+N_FEATURES = len(SELECTED_FEATURES)
+
+# =========================
+# PREDICTION FUNCTION
+# =========================
+def predict_csv(file):
+    df = pd.read_csv(file)
+
+    # Drop label columns if exist
+    df = df.drop(columns=["Label", "label", "class", "Class"], errors="ignore")
+
+    # Check missing features
+    missing_features = [f for f in SELECTED_FEATURES if f not in df.columns]
+    if missing_features:
+        return (
+            f"Missing required features: {missing_features}"
+        )
+
+    # Keep only selected features & correct order
+    feature_df = df[SELECTED_FEATURES].copy()
+
+    # Convert to float
+    X = feature_df.values.astype(float)
+
+    # Scale
+    X_scaled = scaler.transform(X)
+
+    # Predict
+    probs = model.predict(X_scaled).reshape(-1)
+    preds = (probs > 0.5).astype(int)
+
+    # Build output dataframe
+    result = df.copy()
+    result.insert(0, "row_id", range(1, len(df) + 1))
+    result["probability_malware"] = probs
+    result["prediction"] = 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="Malware Detection",
+    description=(
+        "Upload a CSV file containing PE features. "
+    )
+)
+
+demo.launch()