src/ml_surface_classifier.py - master - louintik/ML

src/ml_surface_classifier.py
293 строки10 KB
louisa
first_commit
14 май 2026, 22:04
14 май 2026, 22:0451ac654
"""
Сравнение различных ML методов для классификации поверхностей.
"""
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, ExtraTreesClassifier
from sklearn.svm import SVC
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
import joblib
import warnings

warnings.filterwarnings('ignore')


def load_and_prepare_dataset(dataset_path='../data/processed/surface_dataset.csv'):
    """
    Загрузить датасет и подготовить данные для ML.
    """
    df = pd.read_csv(dataset_path)

    # Исключаем мета-колонки
    exclude_cols = ['file', 'press_type', 'surface_label', 'group', 'prefix_suffix_correlation']
    feature_cols = [col for col in df.columns if col not in exclude_cols]

    X = df[feature_cols]
    y = df['surface_label']

    print(f"Размер датасета: {X.shape}")
    print(f"Фичей: {len(feature_cols)}")
    print(f"Классы: {dict(zip(*np.unique(y, return_counts=True)))}")
    print(f"  0 (hard_surface): {(y == 0).sum()}")
    print(f"  1 (elastic_surface): {(y == 1).sum()}")

    return X, y, feature_cols


def get_models_dict():
    """
    Получить словарь всех моделей для сравнения.
    """
    models = {
        'Random Forest': RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            random_state=42,
            class_weight='balanced'
        ),
        'XGBoost': None,  # Будет добавлена если доступна
        'LightGBM': None,  # Будет добавлена если доступна
        'Extra Trees': ExtraTreesClassifier(
            n_estimators=100,
            max_depth=10,
            random_state=42,
            class_weight='balanced'
        ),
        'SVM': SVC(
            kernel='rbf',
            C=1.0,
            random_state=42,
            class_weight='balanced',
            probability=True
        ),
        'Logistic Regression': LogisticRegression(
            random_state=42,
            max_iter=1000,
            class_weight='balanced'
        ),
        'k-NN': KNeighborsClassifier(
            n_neighbors=5,
            weights='distance'
        ),
        'MLP': MLPClassifier(
            hidden_layer_sizes=(100, 50),
            max_iter=500,
            random_state=42,
            early_stopping=True
        )
    }

    # Пробуем добавить XGBoost
    try:
        from xgboost import XGBClassifier
        models['XGBoost'] = XGBClassifier(
            n_estimators=100,
            max_depth=6,
            random_state=42,
            use_label_encoder=False,
            eval_metric='logloss'
        )
    except ImportError:
        print("⚠ XGBoost не установлен. Пропускаем.")
        models['XGBoost'] = None

    # Пробуем добавить LightGBM
    try:
        from lightgbm import LGBMClassifier
        models['LightGBM'] = LGBMClassifier(
            n_estimators=100,
            max_depth=10,
            random_state=42,
            class_weight='balanced',
            verbose=-1
        )
    except ImportError:
        print("⚠ LightGBM не установлен. Пропускаем.")
        models['LightGBM'] = None

    # Удаляем None значения
    models = {k: v for k, v in models.items() if v is not None}

    return models


def compare_all_models(X, y, test_size=0.2):
    """
    Обучить и сравнить все модели.
    Возвращает DataFrame с результатами и словарь обученных моделей.
    """
    # Заполняем NaN медианным значением
    imputer = SimpleImputer(strategy='median')
    X_imputed = imputer.fit_transform(X)

    # Разделяем на train/test
    X_train, X_test, y_train, y_test = train_test_split(
        X_imputed, y, test_size=test_size, random_state=42, stratify=y
    )

    # Для k-NN и MLP нужно масштабирование
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # Получаем модели
    models = get_models_dict()

    results = []
    trained_models = {}

    print("\n" + "=" * 80)
    print("СРАВНЕНИЕ МОДЕЛЕЙ")
    print("=" * 80)

    for name, model in models.items():
        print(f"\n{'─' * 80}")
        print(f"Модель: {name}")
        print(f"{'─' * 80}")

        # Выбираем данные (масштабированные или нет)
        if name in ['k-NN', 'MLP']:
            X_train_use = X_train_scaled
            X_test_use = X_test_scaled
        else:
            X_train_use = X_train
            X_test_use = X_test

        # Обучаем
        model.fit(X_train_use, y_train)

        # Предсказания
        y_pred = model.predict(X_test_use)

        # Оценка
        acc = accuracy_score(y_test, y_pred)
        cm = confusion_matrix(y_test, y_pred)
        cr = classification_report(y_test, y_pred, target_names=['hard_surface', 'elastic_surface'], output_dict=True)

        print(f"Accuracy: {acc:.3f}")
        print(f"\nConfusion Matrix:")
        print(cm)
        print(f"\nClassification Report:")
        print(classification_report(y_test, y_pred, target_names=['hard_surface', 'elastic_surface']))

        # Для Random Forest и tree-based показываем важность фич
        if hasattr(model, 'feature_importances_'):
            feature_importance = pd.DataFrame({
                'feature': X.columns,
                'importance': model.feature_importances_
            }).sort_values('importance', ascending=False)

            print(f"\nТоп-5 важных фич:")
            print(feature_importance.head(5).to_string(index=False))

        # Сохраняем результаты
        results.append({
            'Model': name,
            'Accuracy': acc,
            'Precision_hard': cr['hard_surface']['precision'],
            'Recall_hard': cr['hard_surface']['recall'],
            'F1_hard': cr['hard_surface']['f1-score'],
            'Precision_elastic': cr['elastic_surface']['precision'],
            'Recall_elastic': cr['elastic_surface']['recall'],
            'F1_elastic': cr['elastic_surface']['f1-score'],
            'Support': len(y_test)
        })

        # Сохраняем обученную модель
        trained_models[name] = {
            'model': model,
            'imputer': imputer,
            'scaler': scaler if name in ['k-NN', 'MLP'] else None,
            'feature_cols': X.columns.tolist()
        }

    # Создаём DataFrame с результатами
    df_results = pd.DataFrame(results)
    df_results = df_results.sort_values('Accuracy', ascending=False)

    return df_results, trained_models, X_test, y_test


def print_comparison_table(df_results):
    """
    Вывести красивую таблицу сравнения.
    """
    print("\n" + "=" * 100)
    print("ИТОГОВАЯ ТАБЛИЦА СРАВНЕНИЯ МОДЕЛЕЙ")
    print("=" * 100)

    # Форматируем для вывода
    df_display = df_results.copy()
    for col in ['Accuracy', 'Precision_hard', 'Recall_hard', 'F1_hard', 'Precision_elastic', 'Recall_elastic', 'F1_elastic']:
        df_display[col] = df_display[col].apply(lambda x: f"{x:.3f}")

    print(df_display.to_string(index=False))

    print("\n" + "=" * 100)
    best_model = df_results.iloc[0]
    print(f"ЛУЧШАЯ МОДЕЛЬ: {best_model['Model']} (Accuracy: {best_model['Accuracy']:.3f})")
    print("=" * 100)


def save_best_model(trained_models, df_results, output_dir='../models'):
    """
    Сохранить лучшую модель по точности.
    """
    best_model_name = df_results.iloc[0]['Model']
    best_model_data = trained_models[best_model_name]

    print(f"\nСохранение лучшей модели: {best_model_name}")

    joblib.dump(best_model_data['model'], f'{output_dir}/surface_model.pkl')
    joblib.dump(best_model_data['imputer'], f'{output_dir}/surface_imputer.pkl')
    joblib.dump(best_model_data['feature_cols'], f'{output_dir}/surface_features.pkl')

    if best_model_data['scaler'] is not None:
        joblib.dump(best_model_data['scaler'], f'{output_dir}/surface_scaler.pkl')

    # Сохраняем имя модели для справки
    with open(f'{output_dir}/best_model_name.txt', 'w') as f:
        f.write(best_model_name)

    print(f"✓ Модель сохранена в: {output_dir}/surface_model.pkl")
    print(f"✓ Imputer сохранен в: {output_dir}/surface_imputer.pkl")
    print(f"✓ Список фич сохранен в: {output_dir}/surface_features.pkl")
    if best_model_data['scaler'] is not None:
        print(f"✓ Scaler сохранен в: {output_dir}/surface_scaler.pkl")

    return best_model_name


if __name__ == '__main__':
    print("=" * 80)
    print("ЗАГРУЗКА ДАТАСЕТА")
    print("=" * 80)

    # Загружаем датасет
    X, y, feature_cols = load_and_prepare_dataset()

    print("\n" + "=" * 80)
    print("ОБУЧЕНИЕ И СРАВНЕНИЕ МОДЕЛЕЙ")
    print("=" * 80)

    # Сравниваем все модели
    df_results, trained_models, X_test, y_test = compare_all_models(X, y)

    # Выводим таблицу сравнения
    print_comparison_table(df_results)

    # Сохраняем лучшую модель
    best_model_name = save_best_model(trained_models, df_results)

    print(f"\n{'=' * 80}")
    print(f"✓ Все модели обучены! Лучшая: {best_model_name}")
    print(f"{'=' * 80}")

    print("\nДля предсказания на новых данных запустите:")
    print("  python predict_surface.py")