train_models.py

import sys
import os
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import pickle
from datetime import datetime

MODELS_FOLDER = 'models'
MIN_SAMPLES = 30
TEST_SIZE = 0.2
RANDOM_STATE = 42

RF_PARAMS = {
    'n_estimators': 100,
    'max_depth': 10,
    'min_samples_split': 5,
    'min_samples_leaf': 2,
    'random_state': RANDOM_STATE,
    'n_jobs': -1
}

def print_header(text):
    print(f"\n{'='*70}")
    print(f"  {text}")
    print('='*70)

def print_section(text):
    print(f"\n{text}")
    print('-' * 50)

def load_and_clean_data(csv_file):
    print_section("LOADING DATA")
    if not os.path.exists(csv_file):
        print(f"File not found: {csv_file}")
        return None
    
    df = pd.read_csv(csv_file)
    print(f"   Total samples: {len(df)}")
    
    print_section("CLEANING DATA")
    sensor_cols = [col for col in df.columns if col.startswith('sensor')]
    if not sensor_cols:
        print("No sensor columns found (sensor1, sensor2, etc.)")
        return None

    initial_count = len(df)
    
    # 1. Ensure numeric and drop NaNs
    df_clean = df.copy()
    for col in sensor_cols:
        df_clean[col] = pd.to_numeric(df_clean[col], errors='coerce')
    df_clean = df_clean.dropna(subset=sensor_cols)
    
    # 2. Handle 'Unknown' objects to avoid training on garbage data
    if 'object_type' in df_clean.columns:
        df_clean = df_clean[df_clean['object_type'] != 'Unknown']
    
    # 3. Handle 'All-1.0' error condition 
    # (1.0 is considered a sensor error or too close)
    # Only remove a row if ALL sensors are <= 1.0 (completely invalid state)
    invalid_all = (df_clean[sensor_cols] <= 1.05).all(axis=1)
    df_clean = df_clean[~invalid_all].copy()
    
    removed = initial_count - len(df_clean)
    print(f"   Removed {removed} invalid rows (NaN, Unknown, or All-Sensors-Error)")
    print(f"   Valid samples: {len(df_clean)}")
    
    if len(df_clean) < MIN_SAMPLES:
        print(f"Need at least {MIN_SAMPLES} valid samples, have {len(df_clean)}")
        return None
    
    return df_clean, sensor_cols

def prepare_features(df, sensor_cols):
    print_section("🔧 PREPARING FEATURES")
    X = df[sensor_cols].values
    
    if 'object_type' not in df.columns:
        print("No 'object_type' column found")
        return None, None, None
    
    y_object = df['object_type'].values
    y_movement = df['movement_type'].values if 'movement_type' in df.columns else None
    
    print(f"   Object types ({len(set(y_object))} classes):")
    for obj_type, count in pd.Series(y_object).value_counts().items():
        print(f"      • {obj_type}: {count} samples")
    
    return X, y_object, y_movement

def train_model(X_train, X_test, y_train, y_test, model_name):
    print_section(f"TRAINING {model_name.upper()}")
    clf = RandomForestClassifier(**RF_PARAMS)
    clf.fit(X_train, y_train)
    
    test_pred = clf.predict(X_test)
    test_acc = accuracy_score(y_test, test_pred)
    print(f"   Testing accuracy: {test_acc:.1%}")
    return clf

def save_models(obj_clf, mov_clf, scaler):
    print_section("SAVING MODELS")
    os.makedirs(MODELS_FOLDER, exist_ok=True)
    
    with open(f'{MODELS_FOLDER}/object_classifier.pkl', 'wb') as f:
        pickle.dump(obj_clf, f)
    if mov_clf:
        with open(f'{MODELS_FOLDER}/movement_classifier.pkl', 'wb') as f:
            pickle.dump(mov_clf, f)
    with open(f'{MODELS_FOLDER}/scaler.pkl', 'wb') as f:
        pickle.dump(scaler, f)
    print(f"   Models saved to: {MODELS_FOLDER}/")

def show_feature_importance(clf, sensor_cols):
    print_section("FEATURE IMPORTANCE")
    importances = clf.feature_importances_
    sorted_idx = np.argsort(importances)[::-1]
    
    for idx in sorted_idx:
        importance = importances[idx]
        bar = '█' * int(importance * 50)
        print(f"      {sensor_cols[idx]:15s} {importance:5.3f} {bar}")

def main():
    print_header("AI MODEL TRAINING PIPELINE")
    
    if len(sys.argv) > 1:
        csv_file = sys.argv[1]
    else:
        csv_files = sorted([f for f in os.listdir('data') if f.endswith('.csv')], reverse=True)
        if csv_files:
            csv_file = f'data/{csv_files[0]}'
            print(f"\n   Using latest CSV: {csv_file}")
        else:
            print("\nNo CSV file specified")
            return

    result = load_and_clean_data(csv_file)
    if result is None: return
    
    df_clean, sensor_cols = result
    X, y_object, y_movement = prepare_features(df_clean, sensor_cols)
    
    print_section("SPLITTING & SCALING")
    X_train, X_test, y_obj_train, y_obj_test = train_test_split(
        X, y_object, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y_object
    )
    
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    
    obj_clf = train_model(X_train_scaled, X_test_scaled, y_obj_train, y_obj_test, "Object Classifier")
    
    mov_clf = None
    if y_movement is not None:
        _, _, y_mov_train, y_mov_test = train_test_split(
            X, y_movement, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y_movement
        )
        mov_clf = train_model(X_train_scaled, X_test_scaled, y_mov_train, y_mov_test, "Movement Classifier")
    
    show_feature_importance(obj_clf, sensor_cols)
    save_models(obj_clf, mov_clf, scaler)
    print_header("TRAINING COMPLETE")

if __name__ == "__main__":
    main()