predict_surface.py - master - louintik/ML

predict_surface.py
487 строк19 KB
louisa
first_commit
14 май 2026, 22:04
14 май 2026, 22:0451ac654
"""
Предсказание типа поверхности для новых файлов.
Поддерживает выбор конкретной модели или сравнение всех моделей.
"""
import os
import sys
import argparse
import pandas as pd
import joblib

# Добавляем src/ в путь для импорта functions.py
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src'))
from functions import process_single_file, load_signal_data, find_all_peaks, group_peaks, extract_key_points
from scipy.signal import find_peaks
import plotly.graph_objects as go
import warnings

warnings.filterwarnings('ignore')

# Доступные модели и их префиксы файлов
AVAILABLE_MODELS = {
    'random_forest': 'Random Forest',
    'extra_trees': 'Extra Trees',
    'svm': 'SVM',
    'logistic_regression': 'Logistic Regression',
    'knn': 'k-NN',
    'mlp': 'MLP',
    'xgboost': 'XGBoost',
    'lightgbm': 'LightGBM',
    'best': 'Best Model'  # Старая лучшая модель
}


def visualize_predictions(filepath, predictions_df, model_name, save_path=None):
    """
    Визуализировать сигнал с метками предсказаний модели.
    """
    # Загружаем сигнал
    signal, time = load_signal_data(filepath)
    
    # Создаём график
    fig = go.Figure()
    
    # Добавляем сигнал
    fig.add_trace(
        go.Scatter(
            x=time,
            y=signal,
            mode='lines',
            marker=dict(color='black'),
            name='Сигнал',
            line=dict(width=1.5)
        )
    )
    
    # Добавляем пики
    peaks_plus = find_peaks(signal, height=0.03)
    peaks_minus = find_peaks(-signal, height=0.03)
    
    fig.add_trace(
        go.Scatter(
            x=time[peaks_plus[0]],
            y=signal[peaks_plus[0]],
            mode='markers',
            marker=dict(color='red', size=6, symbol='circle'),
            name='Положительные пики'
        )
    )
    
    fig.add_trace(
        go.Scatter(
            x=time[peaks_minus[0]],
            y=signal[peaks_minus[0]],
            mode='markers',
            marker=dict(color='blue', size=6, symbol='circle'),
            name='Отрицательные пики'
        )
    )
    
    # Для каждого нажатия добавляем зону и аннотацию с предсказанием
    colors = {
        0: 'rgba(0, 150, 255, 0.2)',   # hard_surface - синий
        1: 'rgba(255, 150, 0, 0.2)'    # elastic_surface - оранжевый
    }
    
    border_colors = {
        0: 'rgba(0, 150, 255, 0.8)',
        1: 'rgba(255, 150, 0, 0.8)'
    }
    
    # Загружаем данные для группировки пиков
    sorted_peak_index, sorted_index_to_height = find_all_peaks(signal)
    group_to_indices = group_peaks(sorted_peak_index)
    group_prefix_suffix = extract_key_points(group_to_indices, sorted_index_to_height)
    
    # Сопоставляем предсказания с группами
    for idx, (_, row) in enumerate(predictions_df.iterrows()):
        group = row['group']
        pred_label = row['predicted_surface']
        confidence = row['confidence']
        
        # Находим границы группы
        if group in group_prefix_suffix:
            pts = group_prefix_suffix[group]
            i1 = pts['first_index']
            i3 = pts['last_index']
            
            # Выделяем зону нажатия
            fig.add_vrect(
                x0=time[i1],
                x1=time[i3],
                fillcolor=colors.get(pred_label, 'rgba(200, 200, 200, 0.2)'),
                opacity=0.3,
                layer="below",
                line_width=2,
                line_color=border_colors.get(pred_label, 'gray'),
            )
            
            # Аннотация с предсказанием
            label_ru = 'Твёрдая' if pred_label == 0 else 'Упругая'
            annotation_text = (
                f"Нажатие #{group}<br>"
                f"Предсказание: {label_ru}<br>"
                f"Уверенность: {confidence:.1f}%"
            )
            
            # Позиция аннотации
            y_max_idx = max(range(i1, i3), key=lambda x: abs(signal[x]))
            
            fig.add_annotation(
                x=time[y_max_idx],
                y=signal[y_max_idx],
                text=annotation_text,
                showarrow=True,
                arrowhead=2,
                arrowsize=1,
                arrowwidth=1,
                arrowcolor='black',
                bgcolor='white',
                bordercolor=border_colors.get(pred_label, 'black'),
                borderwidth=1,
                font=dict(size=10),
                yshift=10
            )
    
    # Обновляем layout
    fig.update_layout(
        title=f"Файл: {os.path.basename(filepath)}<br>Модель: {model_name}",
        xaxis_title="Время (с)",
        yaxis_title="Амплитуда (В)",
        width=1400,
        height=800,
        showlegend=True,
        legend=dict(
            yanchor="top",
            y=0.99,
            xanchor="left",
            x=0.01
        )
    )
    
    # Показываем или сохраняем
    if save_path:
        fig.write_html(save_path)
        print(f"✓ Визуализация сохранена: {save_path}")
    else:
        fig.show()
    
    return fig


def predict_file_single_model(filepath, model, imputer, feature_cols, scaler=None, model_name=None):
    """
    Предсказать тип поверхности для одного файла используя одну модель.
    """
    # Извлекаем фичи
    data = process_single_file(filepath)

    if len(data) == 0:
        return None

    df = pd.DataFrame(data)
    X = df[feature_cols]
    X_imputed = imputer.transform(X)

    # Масштабируем если нужно
    if scaler is not None:
        X_scaled = scaler.transform(X_imputed)
        predictions = model.predict(X_scaled)
        probabilities = model.predict_proba(X_scaled)
    else:
        predictions = model.predict(X_imputed)
        probabilities = model.predict_proba(X_imputed)

    df['predicted_surface'] = predictions
    df['predicted_label'] = df['predicted_surface'].map({
        0: 'hard_surface',
        1: 'elastic_surface'
    })
    df['confidence'] = probabilities.max(axis=1) * 100

    return df


def predict_file_all_models(filepath, models_data):
    """
    Предсказать тип поверхности для одного файла используя ВСЕ модели.
    
    models_data: словарь {имя_модели: {model, imputer, feature_cols, scaler}}
    """
    print(f"\nОбработка: {filepath}")
    print("=" * 100)

    # Извлекаем фичи один раз
    data = process_single_file(filepath)

    if len(data) == 0:
        print("⚠ Нажатия не найдены!")
        return None

    df_base = pd.DataFrame(data)
    print(f"Найдено нажатий: {len(df_base)}")

    # Словарь для хранения результатов по каждой модели
    all_predictions = {}

    print("\n" + "=" * 100)
    print("РЕЗУЛЬТАТЫ ПРЕДСКАЗАНИЯ ПО МОДЕЛЯМ")
    print("=" * 100)

    for model_name, model_info in models_data.items():
        X = df_base[model_info['feature_cols']]
        X_imputed = model_info['imputer'].transform(X)

        # Масштабируем если нужно
        if model_info.get('scaler') is not None:
            X_scaled = model_info['scaler'].transform(X_imputed)
            predictions = model_info['model'].predict(X_scaled)
            probabilities = model_info['model'].predict_proba(X_scaled)
        else:
            predictions = model_info['model'].predict(X_imputed)
            probabilities = model_info['model'].predict_proba(X_imputed)

        all_predictions[model_name] = {
            'predictions': predictions,
            'probabilities': probabilities,
            'labels': [model_info['model'].classes_[p] for p in predictions]
        }

    # Выводим результаты для каждого нажатия
    for idx, row in df_base.iterrows():
        print(f"\n{'─' * 100}")
        print(f"Нажатие #{row['group']} | Длительность: {row['word_duration']:.3f}с | Макс. амплитуда: {row['word_max_amp']:.3f}")
        print(f"{'─' * 100}")

        for model_name, pred_data in all_predictions.items():
            pred_label = str(pred_data['labels'][idx])
            confidence = pred_data['probabilities'][idx].max() * 100
            pred_class = pred_data['predictions'][idx]

            label_ru = 'твердая' if pred_class == 0 else 'упругая'
            print(f"  {model_name:25s} → {pred_label:15s} ({label_ru:8s}) | Уверенность: {confidence:5.1f}%")

    # Общая статистика по моделям
    print("\n" + "=" * 100)
    print("СВОДНАЯ ТАБЛИЦА ПО МОДЕЛЯМ")
    print("=" * 100)

    summary_data = []
    for model_name, pred_data in all_predictions.items():
        counts = pd.Series(pred_data['predictions']).value_counts().to_dict()
        hard_count = counts.get(0, 0)
        elastic_count = counts.get(1, 0)
        avg_conf = pred_data['probabilities'].max(axis=1).mean() * 100

        summary_data.append({
            'Модель': model_name,
            'hard_surface': hard_count,
            'elastic_surface': elastic_count,
            'Средняя уверенность': f"{avg_conf:.1f}%"
        })

    df_summary = pd.DataFrame(summary_data)
    print(df_summary.to_string(index=False))

    # Создаём итоговый DataFrame с предсказаниями лучшей модели
    best_model_name = list(models_data.keys())[0]  # Первая модель из списка
    df_result = df_base.copy()
    df_result['predicted_surface'] = all_predictions[best_model_name]['predictions']
    df_result['predicted_label'] = df_result['predicted_surface'].map({
        0: 'hard_surface (твердая)',
        1: 'elastic_surface (упругая)'
    })
    df_result['confidence'] = all_predictions[best_model_name]['probabilities'].max(axis=1) * 100

    return df_result


def load_model(model_name='best', models_dir='../models'):
    """
    Загрузить конкретную модель по имени.
    
    model_name: одна из 'random_forest', 'extra_trees', 'svm', 'logistic_regression',
                'knn', 'mlp', 'xgboost', 'lightgbm', 'best'
    """
    if model_name not in AVAILABLE_MODELS:
        print(f"❌ Модель '{model_name}' не найдена!")
        print(f"Доступные модели: {', '.join(AVAILABLE_MODELS.keys())}")
        return None
    
    if model_name == 'best':
        # Старая лучшая модель
        model_path = f'{models_dir}/surface_model.pkl'
        imputer_path = f'{models_dir}/surface_imputer.pkl'
        features_path = f'{models_dir}/surface_features.pkl'
        scaler_path = f'{models_dir}/surface_scaler.pkl'
        
        # Читаем имя лучшей модели
        best_model_name_path = f'{models_dir}/best_model_name.txt'
        display_name = 'Best Model'
        if os.path.exists(best_model_name_path):
            with open(best_model_name_path, 'r') as f:
                display_name = f"Best Model ({f.read().strip()})"
    else:
        # Новая отдельная модель
        model_path = f'{models_dir}/{model_name}_model.pkl'
        imputer_path = f'{models_dir}/{model_name}_imputer.pkl'
        features_path = f'{models_dir}/{model_name}_features.pkl'
        scaler_path = f'{models_dir}/{model_name}_scaler.pkl'
        display_name = AVAILABLE_MODELS[model_name]
    
    # Проверяем существование файлов
    if not os.path.exists(model_path):
        print(f"❌ Файл модели не найден: {model_path}")
        print(f"Сначала обучите модель: python train_model.py --model {model_name}")
        return None
    
    if not os.path.exists(imputer_path):
        print(f"❌ Файл импутера не найден: {imputer_path}")
        return None
    
    if not os.path.exists(features_path):
        print(f"❌ Файл фич не найден: {features_path}")
        return None
    
    # Загружаем модель
    model = joblib.load(model_path)
    imputer = joblib.load(imputer_path)
    feature_cols = joblib.load(features_path)
    
    # Загружаем scaler если существует
    scaler = None
    if os.path.exists(scaler_path):
        scaler = joblib.load(scaler_path)
    
    return {
        'model': model,
        'imputer': imputer,
        'feature_cols': feature_cols,
        'scaler': scaler,
        'name': display_name
    }


def load_all_models_for_comparison(models_dir='../models'):
    """
    Загрузить все модели для сравнения.
    """
    models_data = {}
    
    for model_key, model_display_name in AVAILABLE_MODELS.items():
        if model_key == 'best':
            continue  # Пропускаем 'best', загружаем только конкретные модели
        
        model_data = load_model(model_key, models_dir)
        if model_data is not None:
            models_data[model_display_name] = model_data
    
    return models_data


if __name__ == '__main__':
    # Определяем корень проекта (где лежит этот скрипт)
    PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
    
    # Парсим аргументы командной строки
    parser = argparse.ArgumentParser(description='Предсказание типа поверхности')
    parser.add_argument('--model', type=str, default='random_forest',
                        help=f"Модель для предсказания: {', '.join(AVAILABLE_MODELS.keys())}")
    parser.add_argument('--file', type=str,
                        default=os.path.join(PROJECT_ROOT, 'data', 'data_for_check_ml', 'up_20.csv'),
                        help="Путь к файлу для предсказания")
    parser.add_argument('--compare', action='store_true',
                        help="Сравнить все модели")
    parser.add_argument('--models-dir', type=str,
                        default=os.path.join(PROJECT_ROOT, 'models'),
                        help="Директория с моделями")
    parser.add_argument('--no-viz', action='store_true',
                        help="Не создавать визуализацию")

    args = parser.parse_args()
    
    if args.compare:
        # Режим сравнения всех моделей
        print("=" * 100)
        print("ЗАГРУЗКА ВСЕХ МОДЕЛЕЙ ДЛЯ СРАВНЕНИЯ")
        print("=" * 100)

        models_data = load_all_models_for_comparison(args.models_dir)

        if len(models_data) == 0:
            print("❌ Не найдено ни одной обученной модели!")
            print("Сначала обучите модели: python train_model.py --model <name>")
            sys.exit(1)

        print(f"✓ Загружено моделей: {len(models_data)}")
        for name in models_data.keys():
            print(f"  - {name}")

        # Предсказываем для файла
        if os.path.exists(args.file):
            df_result = predict_file_all_models(args.file, models_data)

            if df_result is not None:
                output_path = os.path.join(PROJECT_ROOT, 'outputs', 'prediction_comparison.csv')
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                df_result.to_csv(output_path, index=False)
                print(f"\n✓ Результат сохранён: {output_path}")
        else:
            print(f"⚠ Файл {args.file} не найден!")
    else:
        # Режим одной модели
        print("=" * 100)
        print(f"ЗАГРУЗКА МОДЕЛИ: {AVAILABLE_MODELS.get(args.model, args.model)}")
        print("=" * 100)

        model_data = load_model(args.model, args.models_dir)

        if model_data is None:
            print("❌ Не удалось загрузить модель!")
            sys.exit(1)

        print(f"✓ Модель загружена: {model_data['name']}")

        # Предсказываем для файла
        if os.path.exists(args.file):
            df_result = predict_file_single_model(
                args.file,
                model_data['model'],
                model_data['imputer'],
                model_data['feature_cols'],
                model_data['scaler'],
                model_data['name']
            )

            if df_result is not None:
                print(f"\n{'=' * 100}")
                print("РЕЗУЛЬТАТ ПРЕДСКАЗАНИЯ")
                print(f"{'=' * 100}")

                for idx, row in df_result.iterrows():
                    label_ru = 'твердая' if row['predicted_surface'] == 0 else 'упругая'
                    print(f"Нажатие #{row['group']} | {row['predicted_label']:15s} ({label_ru:8s}) | "
                          f"Уверенность: {row['confidence']:.1f}%")

                output_path = os.path.join(PROJECT_ROOT, 'outputs', 'prediction_result.csv')
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                df_result.to_csv(output_path, index=False)
                print(f"\n✓ Результат сохранён: {output_path}")
                
                # Создаём визуализацию
                if not args.no_viz:
                    from datetime import datetime
                    date_str = datetime.now().strftime('%y_%m_%d_%H%M%S')
                    model_short = model_data['name'].replace(' ', '_').lower()
                    file_base = os.path.splitext(os.path.basename(args.file))[0]
                    viz_filename = f'{file_base}_{model_short}_{date_str}.html'
                    viz_path = os.path.join(PROJECT_ROOT, 'outputs', viz_filename)
                    visualize_predictions(args.file, df_result, model_data['name'], save_path=viz_path)
        else:
            print(f"⚠ Файл {args.file} не найден!")
            print("\nДоступные файлы в data/data_for_check_ml/:")
            check_ml_dir = os.path.join(PROJECT_ROOT, 'data', 'data_for_check_ml')
            if os.path.exists(check_ml_dir):
                for f in os.listdir(check_ml_dir):
                    if f.endswith('.csv'):
                        print(f"  - {f}")