Update app.py

app.py

@@ -4,101 +4,49 @@ 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
+N_FEATURES = model.input_shape[1]
 
-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"
-]
+feature_names = [f"feature_{i+1}" for i in range(N_FEATURES)]
 
-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()
+def predict_malware(*inputs):
+    # inputs → DataFrame
+    X = pd.DataFrame([inputs], columns=feature_names)
 
-    # Convert to float
-    X = feature_df.values.astype(float)
-
-    # Scale
+    # scale
     X_scaled = scaler.transform(X)
 
-    # Predict
-    probs = model.predict(X_scaled).reshape(-1)
-    preds = (probs > 0.5).astype(int)
+    # predict
+    prob = model.predict(X_scaled, verbose=0)[0][0]
+    pred = int(prob >= 0.5)
+
+    label = "Malware" if pred == 1 else "Benign"
+    return label, float(prob)
 
-    # 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"}
-    )
+# UI
+inputs = [
+    gr.Number(label=feat, value=0.0)
+    for feat in feature_names
+]
 
-    return result
+outputs = [
+    gr.Textbox(label="Prediction"),
+    gr.Number(label="Malware Probability")
+]
 
-# =========================
-# 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. "
-    )
+app = gr.Interface(
+    fn=predict_malware,
+    inputs=inputs,
+    outputs=outputs,
+    title="MLP-based Malware Detection",
+    description="Malware detection using MLP neural network + StandardScaler"
 )
 
-demo.launch()
+if __name__ == "__main__":
+    app.launch()