mboukabous/train_unsupervised

1 2"""3train_anomaly_detection.py4 5Trains an anomaly detection model (Isolation Forest, One-Class SVM, etc.) on a dataset.6Allows dropping or selecting columns, label-encoding for non-numeric data,7saves predictions (0 = normal, 1 = outlier) and optionally visualizes in 2D.8 9Usage Example:10--------------11python scripts/train_anomaly_detection.py \12    --model_module isolation_forest \13    --data_path data/raw/my_dataset.csv \14    --drop_columns "unwanted_col" \15    --select_columns "feat1,feat2,feat3" \16    --visualize17"""18 19import os20import sys21import argparse22import importlib23import pandas as pd24import numpy as np25import joblib26 27from sklearn.preprocessing import LabelEncoder28import matplotlib.pyplot as plt29from timeit import default_timer as timer30 31def main(args):32    # Change to the project root if needed33    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))34    os.chdir(project_root)35    sys.path.insert(0, project_root)36 37    # Dynamically import the chosen anomaly model module38    model_module_path = f"models.unsupervised.anomaly.{args.model_module}"39    model_module = importlib.import_module(model_module_path)40 41    # Retrieve the estimator from the model file42    estimator = model_module.estimator43 44    # Prepare results directory45    if args.results_path is None:46        # e.g., 'results/IsolationForest_Anomaly'47        args.results_path = os.path.join("results", f"{estimator.__class__.__name__}_Anomaly")48    os.makedirs(args.results_path, exist_ok=True)49 50    # Prepare model directory51    if args.model_path is None:52        # e.g., 'saved_model/IsolationForest_Anomaly'53        args.model_path = os.path.join('saved_models', f"{estimator.__class__.__name__}_Anomaly")54    os.makedirs(args.model_path, exist_ok=True)55 56    # Load data57    df = pd.read_csv(args.data_path)58    print(f"Data loaded from {args.data_path}, initial shape: {df.shape}")59 60    # Drop empty columns61    df = df.dropna(axis='columns', how='all')62    print("After dropping empty columns:", df.shape)63 64    # Drop specified columns if any65    if args.drop_columns:66        drop_cols = [c.strip() for c in args.drop_columns.split(',') if c.strip()]67        df.drop(columns=drop_cols, inplace=True, errors='ignore')68        print(f"Dropped columns: {drop_cols} | New shape: {df.shape}")69 70    # Select specified columns if any71    if args.select_columns:72        keep_cols = [c.strip() for c in args.select_columns.split(',') if c.strip()]73        df = df[keep_cols]74        print(f"Selected columns: {keep_cols} | New shape: {df.shape}")75 76    # Label-encode non-numeric columns77    for col in df.columns:78        if df[col].dtype == 'object':79            le = LabelEncoder()80            df[col] = le.fit_transform(df[col])81 82    # Convert DataFrame to numpy array83    X = df.values84    print(f"Final data shape after dropping/selecting columns and encoding: {X.shape}")85 86    # Fit the anomaly model87    start_time = timer()88    estimator.fit(X)89    end_time = timer()90    train_time = end_time - start_time91    print(f"Anomaly detection training with {args.model_module} completed in {train_time:.2f} seconds.")92 93    # Save the model94    model_output_path = os.path.join(args.model_path, "anomaly_model.pkl")95    joblib.dump(estimator, model_output_path)96    print(f"Model saved to {model_output_path}")97 98    # Predict outliers: Typically returns 1 for inliers, -1 for outliers (or vice versa)99    # We'll unify them to 0 = normal, 1 = outlier100    raw_preds = estimator.predict(X)101    # Some anomaly detectors do the opposite: IsolationForest => +1 inlier, -1 outlier102    # Convert to 0/1:103    preds_binary = np.where(raw_preds == 1, 0, 1)104 105    outlier_count = np.sum(preds_binary)106    inlier_count = len(preds_binary) - outlier_count107    print(f"Detected {outlier_count} outliers out of {len(X)} samples. ({inlier_count} normal)")108 109    # Save predictions110    pred_df = pd.DataFrame({111        'OutlierPrediction': preds_binary112    })113    pred_path = os.path.join(args.results_path, "predictions.csv")114    pred_df.to_csv(pred_path, index=False)115    print(f"Predictions saved to {pred_path}")116 117    # Visualization if 2D or 3D118    if args.visualize:119        print("Creating anomaly detection visualization...")120        # We'll do PCA => 2D if dimension > 2121        if X.shape[1] > 2:122            from sklearn.decomposition import PCA123            pca = PCA(n_components=2)124            X_2d = pca.fit_transform(X)125            x_label = "PC1"126            y_label = "PC2"127        elif X.shape[1] == 2:128            X_2d = X129            x_label = df.columns[0] if df.shape[1] == 2 else "Feature 1"130            y_label = df.columns[1] if df.shape[1] == 2 else "Feature 2"131        else:132            # 1D or 0D => skip133            print("Only 1 feature or none; can't create 2D scatter. Skipping.")134            return135 136        # Plot137        plt.figure(figsize=(6,5))138        # color outliers differently139        colors = np.where(preds_binary == 1, 'r', 'b')140        plt.scatter(X_2d[:,0], X_2d[:,1], c=colors, s=30, alpha=0.7)141        plt.title(f"{estimator.__class__.__name__} Anomaly Detection")142        plt.xlabel(x_label)143        plt.ylabel(y_label)144 145        # Save146        plot_path = os.path.join(args.results_path, "anomaly_plot.png")147        plt.savefig(plot_path)148        plt.show()149        print(f"Anomaly plot saved to {plot_path}")150 151 152if __name__ == "__main__":153    parser = argparse.ArgumentParser(description="Train an anomaly detection model.")154    parser.add_argument('--model_module', type=str, required=True,155                        help='Name of the anomaly detection model (e.g. isolation_forest, one_class_svm).')156    parser.add_argument('--data_path', type=str, required=True,157                        help='Path to the CSV dataset file.')158    parser.add_argument('--model_path', type=str, default=None,159                        help='Path to save the trained model.')160    parser.add_argument('--results_path', type=str, default=None,161                        help='Directory to save results (predictions, plots).')162    parser.add_argument('--drop_columns', type=str, default='',163                        help='Comma-separated column names to drop.')164    parser.add_argument('--select_columns', type=str, default='',165                        help='Comma-separated column names to keep (ignore the rest).')166    parser.add_argument('--visualize', action='store_true',167                        help='If set, reduce to 2D (via PCA if needed) and color outliers vs. normal points.')168    args = parser.parse_args()169    main(args)170