Update app.py

app.py

@@ -12,19 +12,30 @@ N_FEATURES = model.input_shape[1]
 
 def predict_csv(file):
     df = pd.read_csv(file)
-    # Drop label column
-    df = df.drop(columns=["Label", "label", "Class", "class"], errors="ignore")
-    # Add row index column
+
+    # Drop label column if it exists
+    df = df.drop(columns=["Label", "label"], errors="ignore")
+
+    # Add row index for display only
     df.insert(0, "row_id", range(1, len(df) + 1))
-    # Check number of features
-    if df.shape[1] != N_FEATURES:
-        return f"Expected {N_FEATURES} features, but got {df.shape[1]} columns."
 
-    X = df.drop(columns=["row_id"]).values.astype(float)
+    # Separate features for model
+    feature_df = df.drop(columns=["row_id"])
+
+    # Check feature count
+    if feature_df.shape[1] != N_FEATURES:
+        return (
+            f"Expected {N_FEATURES} features, "
+            f"but got {feature_df.shape[1]} columns."
+        )
+
+    X = feature_df.values.astype(float)
     X_scaled = scaler.transform(X)
+
     probs = model.predict(X_scaled).reshape(-1)
     preds = (probs > 0.5).astype(int)
-    # Build result dataframe
+
+    # Build result table (row_id kept)
     result = df.copy()
     result["probability_malware"] = probs
     result["prediction"] = preds
@@ -35,6 +46,7 @@ def predict_csv(file):
     return result
 
 
+
 demo = gr.Interface(
     fn=predict_csv,
     inputs=gr.File(label="Upload CSV file"),