Update app.py

app.py

@@ -4,57 +4,71 @@ import joblib
 from huggingface_hub import hf_hub_download
 
 # =====================================================
-# Load CICIDS2018 Model + Artifacts (13 classes)
+# Load CICIDS2018 Model
 # =====================================================
 rf_model_path = hf_hub_download(
     repo_id="CodebaseAi/netraids-ml-models",
     filename="rf_pipeline.joblib"
 )
+rf_model = joblib.load(rf_model_path)
 
-artifacts_path = hf_hub_download(
-    repo_id="CodebaseAi/netraids-ml-models",
-    filename="training_artifacts.joblib"
-)
+# =====================================================
+# Load Training Artifacts (SAFE)
+# =====================================================
+artifacts = {}
+try:
+    artifacts_path = hf_hub_download(
+        repo_id="CodebaseAi/netraids-ml-models",
+        filename="training_artifacts.joblib"
+    )
+    artifacts = joblib.load(artifacts_path)
+except Exception as e:
+    print("Artifacts not loaded:", e)
 
-rf_model = joblib.load(rf_model_path)
-artifacts = joblib.load(artifacts_path)
+# Try to infer feature columns safely
+FEATURE_COLUMNS = None
+for key in ["feature_columns", "features", "columns", "X_columns"]:
+    if key in artifacts:
+        FEATURE_COLUMNS = artifacts[key]
+        break
 
-FEATURE_COLUMNS = artifacts["feature_columns"]
-CLASS_MAPPING = artifacts["class_mapping"]
+# Try to infer class mapping safely
+CLASS_MAPPING = None
+for key in ["class_mapping", "label_mapping", "classes", "target_mapping"]:
+    if key in artifacts:
+        CLASS_MAPPING = artifacts[key]
+        break
 
 # =====================================================
-# Load BCC-Darknet Model (5 classes)
+# Load BCC-Darknet Model (5-class)
 # =====================================================
-darknet_model_path = hf_hub_download(
-    repo_id="CodebaseAi/netraids-ml-models",
-    filename="realtime_model.pkl"
+darknet_model = joblib.load(
+    hf_hub_download("CodebaseAi/netraids-ml-models", "realtime_model.pkl")
 )
-darknet_scaler_path = hf_hub_download(
-    repo_id="CodebaseAi/netraids-ml-models",
-    filename="realtime_scaler.pkl"
+darknet_scaler = joblib.load(
+    hf_hub_download("CodebaseAi/netraids-ml-models", "realtime_scaler.pkl")
 )
-darknet_encoder_path = hf_hub_download(
-    repo_id="CodebaseAi/netraids-ml-models",
-    filename="realtime_encoder.pkl"
+darknet_encoder = joblib.load(
+    hf_hub_download("CodebaseAi/netraids-ml-models", "realtime_encoder.pkl")
 )
 
-darknet_model = joblib.load(darknet_model_path)
-darknet_scaler = joblib.load(darknet_scaler_path)
-darknet_encoder = joblib.load(darknet_encoder_path)
-
 # =====================================================
 # Prediction Router
 # =====================================================
 def predict(model_choice, features: dict):
-
     df = pd.DataFrame([features])
 
     if model_choice == "CICIDS2018 (13 Classes)":
-        # Enforce feature order
-        df = df[FEATURE_COLUMNS]
+        # Enforce feature order ONLY if available
+        if FEATURE_COLUMNS is not None:
+            df = df[FEATURE_COLUMNS]
 
         pred_idx = rf_model.predict(df)[0]
-        pred_label = CLASS_MAPPING[pred_idx]
+
+        if CLASS_MAPPING is not None:
+            pred_label = CLASS_MAPPING.get(pred_idx, str(pred_idx))
+        else:
+            pred_label = str(pred_idx)
 
         return {
             "dataset": "CICIDS2018",
@@ -62,7 +76,6 @@ def predict(model_choice, features: dict):
         }
 
     else:
-        # BCC-Darknet pipeline
         X_scaled = darknet_scaler.transform(df)
         pred_encoded = darknet_model.predict(X_scaled)[0]
         pred_label = darknet_encoder.inverse_transform([pred_encoded])[0]
@@ -73,7 +86,7 @@ def predict(model_choice, features: dict):
         }
 
 # =====================================================
-# Gradio Interface (Backend Demo Only)
+# Gradio UI
 # =====================================================
 app = gr.Interface(
     fn=predict,
@@ -89,7 +102,8 @@ app = gr.Interface(
     ],
     outputs=gr.JSON(label="Detection Result"),
     title="NetraIDS – Dual-Model Network Intrusion Detection",
-    description="Cloud-deployed ML inference engine for intrusion detection"
+    description="Robust cloud-deployed ML inference with dataset-specific pipelines"
 )
 
 app.launch()
+