train.py

import os
import glob
import json
import argparse
import numpy as np
import pandas as pd
import joblib
from sklearn.model_selection import StratifiedKFold, train_test_split
import xgboost as xgb
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

from features import extract_features, get_feature_names

def ensure_reports_dir():
    reports_dir = os.path.join(os.path.dirname(__file__), "reports")
    os.makedirs(reports_dir, exist_ok=True)
    return reports_dir

def gather_data(data_dir: str):
    cache_dir = os.path.join(data_dir, "cache")
    os.makedirs(cache_dir, exist_ok=True)
    cache_x = os.path.join(cache_dir, "X.npy")
    cache_y = os.path.join(cache_dir, "y.npy")
    cache_g = os.path.join(cache_dir, "groups.npy")
    
    if os.path.exists(cache_x) and os.path.exists(cache_y) and os.path.exists(cache_g):
        print("Loading cached features from disk...")
        return np.load(cache_x), np.load(cache_y), np.load(cache_g)

    X = []
    y = []
    groups = []
    
    real_dir = os.path.join(data_dir, "real")
    screen_dir = os.path.join(data_dir, "screen")
    
    import cv2
    from tqdm import tqdm
    import re
    
    print(f"Scanning for images in {real_dir} and {screen_dir}...")
    
    def extract_group(filepath, is_screen):
        basename = os.path.basename(filepath)
        if not is_screen:
            match = re.search(r'(?:DS-05-)(\d+)', basename)
            return match.group(1) if match else basename
        else:
            match = re.search(r'-(\d{3,4})\.\w+$', basename)
            return match.group(1) if match else basename

    def process_image(filepath, label, is_screen):
        img = cv2.imread(filepath)
        if img is not None:
            vec, _ = extract_features(img)
            grp = extract_group(filepath, is_screen)
            return vec, label, grp
        return None

    from joblib import Parallel, delayed

    real_files = glob.glob(os.path.join(real_dir, "*.*"))
    screen_files = glob.glob(os.path.join(screen_dir, "*.*"))
    
    print(f"Extracting features from {len(real_files)} real and {len(screen_files)} screen images in parallel (n_jobs=4 to save RAM)...")
    
    real_results = Parallel(n_jobs=4)(delayed(process_image)(f, 0, False) for f in tqdm(real_files, desc="Real Images"))
    screen_results = Parallel(n_jobs=4)(delayed(process_image)(f, 1, True) for f in tqdm(screen_files, desc="Screen Images"))
    
    all_results = [r for r in real_results + screen_results if r is not None]
    
    for vec, label, grp in all_results:
        X.append(vec)
        y.append(label)
        groups.append(grp)
            
    X_arr = np.array(X)
    y_arr = np.array(y)
    g_arr = np.array(groups)
    
    print("Saving features to cache...")
    np.save(cache_x, X_arr)
    np.save(cache_y, y_arr)
    np.save(cache_g, g_arr)
    
    return X_arr, y_arr, g_arr

def train_model(data_dir: str, phone_data_dir: str = None, mode: str = "default"):
    print("Starting training phase...")
    X, y, groups = gather_data(data_dir)
    
    if len(X) == 0:
        print("Error: No valid images found in dataset directories.")
        return
        
    print(f"Successfully extracted features from {len(X)} images.")
    
    if len(np.unique(y)) < 2:
        print("Error: Dataset must contain both real and screen images.")
        return

    from sklearn.model_selection import GroupShuffleSplit
    
    # --- Grouped Split (Leakage Free) ---
    gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
    train_idx, test_idx = next(gss.split(X, y, groups))
    X_tr_g, X_te_g, y_tr_g, y_te_g = X[train_idx], X[test_idx], y[train_idx], y[test_idx]
    from sklearn.calibration import CalibratedClassifierCV
    
    # Base robust Random Forest to compare
    ratio_g = float(np.sum(y_tr_g == 0)) / np.sum(y_tr_g == 1) if np.sum(y_tr_g == 1) > 0 else 1.0
    rf_base = xgb.XGBClassifier(n_estimators=100, max_depth=4, reg_lambda=10, random_state=42, scale_pos_weight=ratio_g, tree_method='hist', device='cuda')
    
    # Wrap in Logistic Calibration (Platt Scaling)
    rf_grouped = CalibratedClassifierCV(rf_base, method='sigmoid', cv=3)
    
    print("Training Calibrated XGBoost model on GPU (Leakage-Free)...")
    rf_grouped.fit(X_tr_g, y_tr_g)
    
    print("Evaluating leak-free model on test set...")
    y_pred = rf_grouped.predict(X_te_g)
    y_proba = rf_grouped.predict_proba(X_te_g)[:, 1]
    
    acc = accuracy_score(y_te_g, y_pred)
    prec = precision_score(y_te_g, y_pred, zero_division=0)
    rec = recall_score(y_te_g, y_pred, zero_division=0)
    f1 = f1_score(y_te_g, y_pred, zero_division=0)
    
    try:
        roc_auc = roc_auc_score(y_te_g, y_proba)
    except ValueError:
        roc_auc = 0.0

    cm = confusion_matrix(y_te_g, y_pred)

    print("------------------------------")
    print("Metrics on Grouped Held-out Test Set:")
    print(f"Accuracy:  {acc:.4f}")
    print(f"Precision: {prec:.4f}")
    print(f"Recall:    {rec:.4f}")
    print(f"F1 Score:  {f1:.4f}")
    print(f"ROC AUC:   {roc_auc:.4f}")
    print("Confusion Matrix:")
    print(cm)
    print("------------------------------")

    model_path = os.path.join(os.path.dirname(__file__), "model.joblib")
    joblib.dump(rf_grouped, model_path)
    print(f"Model saved to {model_path}.")

    # --- Phone Domain Calibration & Model Comparison ---
    threshold = 0.50
    calibration_status = "Uncalibrated"
    final_model = rf_grouped
    final_model_name = "XGBoost Classifier"
    
    if mode in ["hybrid-domain", "phone-adapted"] and phone_data_dir:
        print(f"Running phone-domain model comparison using {phone_data_dir}...")
        X_phone, y_phone, _ = gather_data(phone_data_dir)
        if len(X_phone) > 0:
            from sklearn.model_selection import StratifiedKFold
            skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
            
            # Tracking metrics for the 4 models
            f1_icl_base = []
            f1_thresh_cal = []
            f1_phone_adapt = []
            f1_rule_boost = []
            
            f_names = get_feature_names()
            idx_bezel = f_names.index("bezel_score")
            idx_persp = f_names.index("perspective_score")
            idx_glare = f_names.index("glare_patch_size")

            best_threshold_overall = 0.50
            best_f1_overall = 0.0

            for train_idx_p, test_idx_p in skf.split(X_phone, y_phone):
                X_tr_p, y_tr_p = X_phone[train_idx_p], y_phone[train_idx_p]
                X_te_p, y_te_p = X_phone[test_idx_p], y_phone[test_idx_p]
                
                # 1. ICL Base Model
                probs_base = rf_grouped.predict_proba(X_te_p)[:, 1]
                preds_base = (probs_base >= 0.5).astype(int)
                f1_icl_base.append(f1_score(y_te_p, preds_base, zero_division=0))
                
                # 2. Threshold-Calibrated
                best_th = 0.5
                best_th_f1 = 0
                probs_tr_p = rf_grouped.predict_proba(X_tr_p)[:, 1]
                for th in np.arange(0.1, 0.9, 0.05):
                    f1 = f1_score(y_tr_p, (probs_tr_p >= th).astype(int), zero_division=0)
                    if f1 > best_th_f1:
                        best_th_f1 = f1
                        best_th = th
                
                preds_th = (probs_base >= best_th).astype(int)
                f1_thresh_cal.append(f1_score(y_te_p, preds_th, zero_division=0))
                
                # Update best overall threshold (averaging over folds)
                best_threshold_overall += best_th / 5.0
                
                # 3. Rule-Boosted Hybrid (Applied to Base probs)
                boosted_probs = probs_base.copy()
                for i in range(len(X_te_p)):
                    boost = 0.0
                    if X_te_p[i, idx_bezel] > 0.02: boost += 0.15
                    if X_te_p[i, idx_persp] > 0.2: boost += 0.15
                    if X_te_p[i, idx_glare] > 0.02: boost += 0.1
                    boosted_probs[i] = min(1.0, boosted_probs[i] + boost)
                
                preds_boost = (boosted_probs >= best_th).astype(int)
                f1_rule_boost.append(f1_score(y_te_p, preds_boost, zero_division=0))
                
                # 4. Phone-adapted Model (Train on ICL + Phone Train fold)
                X_comb = np.vstack([X_tr_g, X_tr_p])
                y_comb = np.concatenate([y_tr_g, y_tr_p])
                ratio_c = float(np.sum(y_comb == 0)) / np.sum(y_comb == 1) if np.sum(y_comb == 1) > 0 else 1.0
                rf_adapted_base = xgb.XGBClassifier(n_estimators=100, max_depth=4, reg_lambda=10, random_state=42, scale_pos_weight=ratio_c, tree_method='hist', device='cuda')
                rf_adapted = CalibratedClassifierCV(rf_adapted_base, method='sigmoid', cv=3)
                rf_adapted.fit(X_comb, y_comb)
                
                probs_adapt = rf_adapted.predict_proba(X_te_p)[:, 1]
                # Also find best threshold for adapted model on its own train set
                probs_adapt_tr = rf_adapted.predict_proba(X_tr_p)[:, 1]
                best_th_adapt = 0.5
                best_th_adapt_f1 = 0
                for th in np.arange(0.1, 0.9, 0.05):
                    f1 = f1_score(y_tr_p, (probs_adapt_tr >= th).astype(int), zero_division=0)
                    if f1 > best_th_adapt_f1:
                        best_th_adapt_f1 = f1
                        best_th_adapt = th
                
                preds_adapt = (probs_adapt >= best_th_adapt).astype(int)
                f1_phone_adapt.append(f1_score(y_te_p, preds_adapt, zero_division=0))

            mean_f1_base = np.mean(f1_icl_base)
            mean_f1_thresh = np.mean(f1_thresh_cal)
            mean_f1_boost = np.mean(f1_rule_boost)
            mean_f1_adapt = np.mean(f1_phone_adapt)

            print("\n=== Phone-Domain 5-Fold CV Model Comparison ===")
            print(f"1. ICL Base Model:                F1 = {mean_f1_base:.4f}")
            print(f"2. Threshold-Calibrated Model:    F1 = {mean_f1_thresh:.4f}")
            print(f"3. Phone-Adapted Model:           F1 = {mean_f1_adapt:.4f}")
            print(f"4. Rule-Boosted Hybrid Model:     F1 = {mean_f1_boost:.4f}")
            print("===============================================\n")

            reports_dir = ensure_reports_dir()
            comp_df = pd.DataFrame({
                "Model": ["ICL Base", "Threshold-Calibrated", "Phone-Adapted", "Rule-Boosted Hybrid"],
                "Phone_CV_F1_Score": [mean_f1_base, mean_f1_thresh, mean_f1_adapt, mean_f1_boost]
            })
            comp_df.to_csv(os.path.join(reports_dir, "model_comparison.csv"), index=False)
            
            with open(os.path.join(reports_dir, "model_selection.md"), "w") as f:
                f.write("# Model Selection Report\n\n")
                f.write("Based on 5-Fold Stratified CV on the phone dataset:\n")
                f.write(comp_df.to_markdown(index=False))

            # Select final model logic
            best_model_idx = np.argmax([mean_f1_base, mean_f1_thresh, mean_f1_adapt, mean_f1_boost])
            if best_model_idx == 2:
                # Retrain phone-adapted on FULL phone set + ICL
                print("Selected: Phone-Adapted Model. Training final model on all ICL + Phone data...")
                X_comb_final = np.vstack([X_tr_g, X_phone])
                y_comb_final = np.concatenate([y_tr_g, y_phone])
                ratio_c = float(np.sum(y_comb_final == 0)) / np.sum(y_comb_final == 1) if np.sum(y_comb_final == 1) > 0 else 1.0
                rf_final_base = xgb.XGBClassifier(n_estimators=100, max_depth=4, reg_lambda=10, random_state=42, scale_pos_weight=ratio_c, tree_method='hist', device='cuda')
                rf_final = CalibratedClassifierCV(rf_final_base, method='sigmoid', cv=3)
                rf_final.fit(X_comb_final, y_comb_final)
                
                # Get best threshold for full phone set
                probs_final = rf_final.predict_proba(X_phone)[:, 1]
                best_th_final = 0.5
                best_th_f1_final = 0
                for th in np.arange(0.1, 0.9, 0.05):
                    f1 = f1_score(y_phone, (probs_final >= th).astype(int), zero_division=0)
                    if f1 > best_th_f1_final:
                        best_th_f1_final = f1
                        best_th_final = th
                        
                final_model = rf_final
                threshold = best_th_final
                calibration_status = "Phone-Adapted (Full Dataset)"
                final_model_name = "Phone-Adapted XGBoost"
                
            elif best_model_idx in [1, 3]:
                print("Selected: Threshold-Calibrated or Rule-Boosted. Using ICL Base with shifted threshold.")
                threshold = best_threshold_overall
                calibration_status = "Phone Domain Calibrated + Boosts (5-Fold CV)"
            else:
                print("Selected: ICL Base. (Unlikely, but fallback).")
                threshold = 0.50
                calibration_status = "Uncalibrated Base"

            joblib.dump(final_model, model_path)
            print(f"Final model saved. Threshold set to: {threshold:.2f}")

    # Update metadata
    meta_path = os.path.join(os.path.dirname(__file__), "model_metadata.json")
    meta = {
        "model_type": final_model_name,
        "threshold": threshold,
        "calibration_status": calibration_status,
        "icl_metrics": {
            "accuracy": float(acc),
            "precision": float(prec),
            "recall": float(rec),
            "f1_score": float(f1),
            "roc_auc": float(roc_auc)
        },
        "feature_count": len(get_feature_names()),
        "feature_names": get_feature_names()
    }
    with open(meta_path, 'w') as f:
        json.dump(meta, f, indent=4)
    print("Updated model_metadata.json.")

    # Save Feature Importance Report if available
    reports_dir = ensure_reports_dir()
    importances = None
    if hasattr(final_model, 'feature_importances_'):
        importances = final_model.feature_importances_
    elif hasattr(final_model, 'estimator') and hasattr(final_model.estimator, 'feature_importances_'):
        importances = final_model.estimator.feature_importances_
    
    if importances is not None:
        fi_df = pd.DataFrame({
            "Feature": get_feature_names(),
            "Importance": importances
        }).sort_values(by="Importance", ascending=False)
        fi_df.to_csv(os.path.join(reports_dir, "feature_importance.csv"), index=False)
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train Fraud Detection Model")
    parser.add_argument("--data", type=str, default="dataset", help="Directory containing real/ and screen/ subfolders")
    parser.add_argument("--phone-data", type=str, default=None, help="Directory containing phone images")
    parser.add_argument("--mode", type=str, default="default", help="Training mode (e.g. hybrid-domain)")
    args = parser.parse_args()
    
    train_model(args.data, args.phone_data, args.mode)