""" Traffic Gate ML Training Script ================================ Trains two models from traffic_gate_dataset.csv: 1. Classifier → predict winner (A or B) 2. Regressor → predict gate_open_seconds Usage: python train_traffic_gate.py python train_traffic_gate.py --data path/to/dataset.csv python train_traffic_gate.py --model xgboost # or randomforest / both (default) """ import argparse import os import warnings import joblib import numpy as np import pandas as pd from sklearn.ensemble import GradientBoostingClassifier, GradientBoostingRegressor from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor from sklearn.metrics import (accuracy_score, classification_report, mean_absolute_error, r2_score) from sklearn.model_selection import cross_val_score, train_test_split from sklearn.pipeline import Pipeline from sklearn.preprocessing import StandardScaler warnings.filterwarnings("ignore") try: from xgboost import XGBClassifier, XGBRegressor XGBOOST_AVAILABLE = True except ImportError: XGBOOST_AVAILABLE = False # ───────────────────────────────────────────── # CONFIG # ───────────────────────────────────────────── FEATURE_COLS = [ "hour", "road_length_a_m", "road_length_b_m", "lane_a_car", "lane_a_motorcycle", "lane_a_bus", "lane_a_truck", "lane_a_total", "lane_b_car", "lane_b_motorcycle", "lane_b_bus", "lane_b_truck", "lane_b_total", # Engineered features (added below) "diff_total", "diff_weighted_score", "road_length_ratio", "rush_hour", ] CLF_TARGET = "winner" REG_TARGET = "gate_open_seconds" OUTPUT_DIR = "models" os.makedirs(OUTPUT_DIR, exist_ok=True) # ───────────────────────────────────────────── # FEATURE ENGINEERING # ───────────────────────────────────────────── WEIGHTS = {"car": 2.0, "motorcycle": 1.5, "bus": 3.0, "truck": 3.5} def engineer_features(df: pd.DataFrame) -> pd.DataFrame: df = df.copy() # Weighted traffic score per lane (mirrors gate time logic) df["score_a"] = ( df["lane_a_car"] * WEIGHTS["car"] + df["lane_a_motorcycle"] * WEIGHTS["motorcycle"] + df["lane_a_bus"] * WEIGHTS["bus"] + df["lane_a_truck"] * WEIGHTS["truck"] ) df["score_b"] = ( df["lane_b_car"] * WEIGHTS["car"] + df["lane_b_motorcycle"] * WEIGHTS["motorcycle"] + df["lane_b_bus"] * WEIGHTS["bus"] + df["lane_b_truck"] * WEIGHTS["truck"] ) # Key delta features the model can learn from df["diff_total"] = df["lane_a_total"] - df["lane_b_total"] df["diff_weighted_score"] = df["score_a"] - df["score_b"] df["road_length_ratio"] = df["road_length_a_m"] / (df["road_length_b_m"] + 1e-9) # Rush hour flag df["rush_hour"] = df["hour"].apply( lambda h: 1 if (7 <= h <= 9 or 17 <= h <= 19) else 0 ) return df # ───────────────────────────────────────────── # LOAD & PREPARE DATA # ───────────────────────────────────────────── def load_data(path: str): df = pd.read_csv(path) df = engineer_features(df) X = df[FEATURE_COLS] y_clf = (df[CLF_TARGET] == "B").astype(int) # 0 = A, 1 = B y_reg = df[REG_TARGET] return X, y_clf, y_reg # ───────────────────────────────────────────── # MODEL DEFINITIONS # ───────────────────────────────────────────── def get_models(model_type: str): clf_models, reg_models = {}, {} if model_type in ("randomforest", "both"): clf_models["RandomForest"] = Pipeline([ ("scaler", StandardScaler()), ("clf", RandomForestClassifier( n_estimators=300, max_depth=8, min_samples_leaf=2, random_state=42, n_jobs=-1)) ]) reg_models["RandomForest"] = Pipeline([ ("scaler", StandardScaler()), ("reg", RandomForestRegressor( n_estimators=300, max_depth=8, min_samples_leaf=2, random_state=42, n_jobs=-1)) ]) if model_type in ("xgboost", "both"): if not XGBOOST_AVAILABLE: print("⚠ XGBoost not installed. Run: pip install xgboost") else: clf_models["XGBoost"] = XGBClassifier( n_estimators=400, max_depth=5, learning_rate=0.05, subsample=0.8, colsample_bytree=0.8, use_label_encoder=False, eval_metric="logloss", random_state=42, n_jobs=-1, verbosity=0) reg_models["XGBoost"] = XGBRegressor( n_estimators=400, max_depth=5, learning_rate=0.05, subsample=0.8, colsample_bytree=0.8, random_state=42, n_jobs=-1, verbosity=0) return clf_models, reg_models # ───────────────────────────────────────────── # TRAINING & EVALUATION # ───────────────────────────────────────────── def train_and_evaluate(X, y_clf, y_reg, clf_models, reg_models): X_train, X_test, yc_train, yc_test, yr_train, yr_test = train_test_split( X, y_clf, y_reg, test_size=0.2, random_state=42, stratify=y_clf ) best_clf, best_reg = None, None best_clf_acc, best_reg_mae = 0, float("inf") print("\n" + "═" * 60) print(" CLASSIFIER → Predict winner (A or B)") print("═" * 60) for name, model in clf_models.items(): model.fit(X_train, yc_train) preds = model.predict(X_test) acc = accuracy_score(yc_test, preds) cv = cross_val_score(model, X, y_clf, cv=5, scoring="accuracy") print(f"\n▶ {name}") print(f" Test Accuracy : {acc:.4f} ({acc*100:.1f}%)") print(f" CV Accuracy : {cv.mean():.4f} ± {cv.std():.4f}") print(classification_report(yc_test, preds, target_names=["Lane A", "Lane B"], digits=3)) if acc > best_clf_acc: best_clf_acc = acc best_clf = (name, model) print("\n" + "═" * 60) print(" REGRESSOR → Predict gate_open_seconds") print("═" * 60) for name, model in reg_models.items(): model.fit(X_train, yr_train) preds = model.predict(X_test) mae = mean_absolute_error(yr_test, preds) r2 = r2_score(yr_test, preds) cv = cross_val_score(model, X, y_reg, cv=5, scoring="neg_mean_absolute_error") print(f"\n▶ {name}") print(f" Test MAE : {mae:.4f} seconds") print(f" Test R² : {r2:.4f}") print(f" CV MAE : {-cv.mean():.4f} ± {cv.std():.4f}") if mae < best_reg_mae: best_reg_mae = mae best_reg = (name, model) return best_clf, best_reg # ───────────────────────────────────────────── # SAVE MODELS # ───────────────────────────────────────────── def save_models(best_clf, best_reg): clf_path = os.path.join(OUTPUT_DIR, "winner_classifier.pkl") reg_path = os.path.join(OUTPUT_DIR, "gate_seconds_regressor.pkl") joblib.dump(best_clf[1], clf_path) joblib.dump(best_reg[1], reg_path) print("\n" + "═" * 60) print(" SAVED MODELS") print("═" * 60) print(f" ✔ Classifier [{best_clf[0]}] → {clf_path}") print(f" ✔ Regressor [{best_reg[0]}] → {reg_path}") return clf_path, reg_path # ───────────────────────────────────────────── # INFERENCE EXAMPLE # ───────────────────────────────────────────── def run_inference_example(clf_path: str, reg_path: str): clf = joblib.load(clf_path) reg = joblib.load(reg_path) sample = pd.DataFrame([{ "hour": 8, "road_length_a_m": 100, "road_length_b_m": 200, "lane_a_car": 1, "lane_a_motorcycle": 0, "lane_a_bus": 0, "lane_a_truck": 0, "lane_a_total": 1, "lane_b_car": 1, "lane_b_motorcycle": 0, "lane_b_bus": 3, "lane_b_truck": 0, "lane_b_total": 4, }]) sample = engineer_features(sample) sample = sample[FEATURE_COLS] winner = "B" if clf.predict(sample)[0] == 1 else "A" gate_secs = reg.predict(sample)[0] print("\n" + "═" * 60) print(" INFERENCE EXAMPLE (matches your original JSON)") print("═" * 60) print(f" Predicted winner : Lane {winner}") print(f" Predicted gate open : {gate_secs:.2f} seconds") print("═" * 60) # ───────────────────────────────────────────── # MAIN # ───────────────────────────────────────────── def main(): parser = argparse.ArgumentParser(description="Train Traffic Gate Model") parser.add_argument("--data", default="traffic_gate_dataset.csv", help="Path to dataset CSV") parser.add_argument("--model", default="both", choices=["randomforest", "xgboost", "both"], help="Model type to train") args = parser.parse_args() print(f"\n📂 Loading data from: {args.data}") X, y_clf, y_reg = load_data(args.data) print(f" Rows: {len(X)} | Features: {len(FEATURE_COLS)}") print(f" Winner distribution: A={( y_clf==0).sum()} B={(y_clf==1).sum()}") clf_models, reg_models = get_models(args.model) best_clf, best_reg = train_and_evaluate(X, y_clf, y_reg, clf_models, reg_models) clf_path, reg_path = save_models(best_clf, best_reg) run_inference_example(clf_path, reg_path) if __name__ == "__main__": main()