Spaces:
Sleeping
Sleeping
| """ | |
| preprocess.py - TemporalDrift-ETM | |
| ===================================== | |
| Converts any NTLFlowLyzer-compatible CSV into the fixed 35-feature | |
| scaled input that the ensemble model expects. | |
| The model was trained on 35 features selected by Random Forest importance | |
| ranking from an original 76-column feature space. Uploaded CSV files may | |
| have any number of columns (40, 50, 70, 85, ...). This module handles: | |
| - Any number of input columns | |
| - Columns that are metadata, identifiers, or labels -> silently skipped | |
| - Selected features present in the CSV -> cleaned and scaled | |
| - Selected features absent from the CSV -> filled with 0.0 | |
| - Non-numeric cell values -> coerced to column mean | |
| - Inf / -Inf values -> replaced with column mean | |
| The per-feature scaling parameters (data_min, scale) are extracted from | |
| the training MinMaxScaler at notebook-save time and stored in scaler_35.pkl, | |
| so this module never requires the full 76-feature scaler at inference time. | |
| Standalone usage | |
| ---------------- | |
| python preprocess.py input.csv output.csv [--models models/] | |
| Module usage (from app.py) | |
| -------------------------- | |
| from preprocess import load_scaler_35, prepare_features | |
| scaler_35 = load_scaler_35("models/scaler_35.pkl") | |
| X, present, miss = prepare_features(df, scaler_35) | |
| """ | |
| import os | |
| import sys | |
| import numpy as np | |
| import pandas as pd | |
| import joblib | |
| # --------------------------------------------------------------------------- | |
| # Column names that are never network-flow features. | |
| # Any CSV column in this set is silently ignored during feature extraction. | |
| # --------------------------------------------------------------------------- | |
| NON_FEATURE_COLS = frozenset({ | |
| "Flow ID", | |
| "Src IP", | |
| "Src Port", | |
| "Dst IP", | |
| "Dst Port", | |
| "Protocol", | |
| "Timestamp", | |
| "Malware Family", | |
| "Label", | |
| "label", | |
| "session", | |
| "Session", | |
| "Prediction", | |
| }) | |
| # --------------------------------------------------------------------------- | |
| # Public API | |
| # --------------------------------------------------------------------------- | |
| def load_scaler_35(path): | |
| """ | |
| Load the 35-feature scaler dict produced by the notebook save cell. | |
| The dict contains three keys: | |
| feature_names : list[str] - the 35 selected feature names (in order) | |
| data_min : np.ndarray - training minimum per feature, shape (35,) | |
| scale : np.ndarray - 1/(max-min) per feature, shape (35,) | |
| Parameters | |
| ---------- | |
| path : str | |
| Path to scaler_35.pkl. | |
| Returns | |
| ------- | |
| dict | |
| """ | |
| artifact = joblib.load(path) | |
| required = {"feature_names", "data_min", "scale"} | |
| missing = required - set(artifact.keys()) | |
| if missing: | |
| raise ValueError( | |
| f"scaler_35.pkl is missing keys: {missing}. " | |
| "Re-run the updated notebook save cell to regenerate it." | |
| ) | |
| return artifact | |
| def prepare_features(df, scaler_35): | |
| """ | |
| Convert a DataFrame with any columns into the 35-feature scaled matrix. | |
| For each of the 35 model features: | |
| - Column present in df -> coerce to float, replace non-finite values | |
| with the column mean, apply MinMaxScaler | |
| formula: (x - data_min) * scale, clip to [0,1]. | |
| - Column absent from df -> fill the entire column with 0.0 in scaled | |
| space (neutral value that does not push the | |
| model toward any particular prediction). | |
| Parameters | |
| ---------- | |
| df : pd.DataFrame | |
| Raw input data (any number of columns, any column names). | |
| scaler_35 : dict | |
| Output of load_scaler_35(). | |
| Returns | |
| ------- | |
| X_scaled : np.ndarray, shape (n_samples, 35) | |
| Model-ready scaled feature matrix. | |
| present : list[str] | |
| Names of the 35 features that were found in df. | |
| missing : list[str] | |
| Names of the 35 features that were absent in df (filled with 0.0). | |
| """ | |
| feat_names = scaler_35["feature_names"] | |
| data_min = np.asarray(scaler_35["data_min"], dtype=np.float64) | |
| scale = np.asarray(scaler_35["scale"], dtype=np.float64) | |
| n = len(df) | |
| X = np.zeros((n, len(feat_names)), dtype=np.float64) | |
| present = [] | |
| missing = [] | |
| for i, feat in enumerate(feat_names): | |
| if feat not in df.columns: | |
| missing.append(feat) | |
| # Column stays at 0.0 - neutral fill in scaled space. | |
| continue | |
| col = pd.to_numeric(df[feat], errors="coerce") | |
| # Replace non-finite values with the finite column mean. | |
| finite_vals = col[np.isfinite(col)] | |
| col_mean = float(finite_vals.mean()) if len(finite_vals) > 0 else 0.0 | |
| col = col.fillna(col_mean) | |
| col = col.replace([np.inf, -np.inf], col_mean) | |
| # MinMaxScaler transform: X_scaled = (X - data_min) * scale | |
| # This exactly replicates sklearn's MinMaxScaler.transform() | |
| # but for a single feature at a time. | |
| X[:, i] = (col.to_numpy(dtype=np.float64) - data_min[i]) * scale[i] | |
| present.append(feat) | |
| # Clip to [0, 1] as sklearn's MinMaxScaler.transform() does. | |
| np.clip(X, 0.0, 1.0, out=X) | |
| return X, present, missing | |
| def report_coverage(present, missing, file=None): | |
| """ | |
| Print a human-readable feature coverage summary. | |
| Parameters | |
| ---------- | |
| present : list[str] | |
| missing : list[str] | |
| file : file object or None (default: stdout) | |
| """ | |
| total = len(present) + len(missing) | |
| print(f"Feature coverage: {len(present)}/{total} columns found in CSV.", file=file) | |
| if missing: | |
| print(f" {len(missing)} feature(s) absent - filled with 0 in scaled space:", | |
| file=file) | |
| for m in missing: | |
| print(f" - {m}", file=file) | |
| # --------------------------------------------------------------------------- | |
| # Standalone CLI | |
| # --------------------------------------------------------------------------- | |
| def _main(): | |
| import argparse | |
| parser = argparse.ArgumentParser( | |
| description=( | |
| "Convert any NTLFlowLyzer-compatible CSV to the 35-feature " | |
| "scaled format expected by TemporalDrift-ETM." | |
| ) | |
| ) | |
| parser.add_argument("input", help="Path to the raw input CSV file") | |
| parser.add_argument("output", help="Path to write the processed output CSV") | |
| parser.add_argument( | |
| "--models", default="models", | |
| help="Directory containing scaler_35.pkl (default: models/)", | |
| ) | |
| args = parser.parse_args() | |
| scaler_35_path = os.path.join(args.models, "scaler_35.pkl") | |
| if not os.path.exists(scaler_35_path): | |
| print( | |
| f"ERROR: {scaler_35_path} not found.\n" | |
| "Re-run the updated notebook save cell to generate scaler_35.pkl.", | |
| file=sys.stderr, | |
| ) | |
| sys.exit(1) | |
| print(f"Reading {args.input} ...") | |
| df = pd.read_csv(args.input) | |
| print(f" {len(df):,} rows | {df.shape[1]} columns") | |
| scaler_35 = load_scaler_35(scaler_35_path) | |
| feat_names = scaler_35["feature_names"] | |
| print(f"\nPreparing {len(feat_names)} model features ...") | |
| X, present, missing = prepare_features(df, scaler_35) | |
| report_coverage(present, missing) | |
| out_df = pd.DataFrame(X, columns=feat_names) | |
| out_df.to_csv(args.output, index=False) | |
| print(f"\nSaved: {args.output}") | |
| print(f"Output shape: {out_df.shape}") | |
| if __name__ == "__main__": | |
| _main() | |