# === Вот эта связка нужна для всех исполняемых скриптов внутри библиотеки, которая будет использоваться как сабмодуль === import sys from pathlib import Path ROOT = Path(__file__).parent.parent if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) # === перед вмеми импортами === import pandas as pd import numpy as np import joblib import json import argparse from sklearn.linear_model import LogisticRegression from sklearn.neural_network import MLPClassifier from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split from sklearn.metrics import classification_report, confusion_matrix, accuracy_score def balance_data(df, target_classes, random_state=42): """ Балансирует только указанные классы до минимального размера среди них. Класс 'none' (и любые другие) остаются без изменений. """ # Разделяем на целевые и остальные target_df = df[df['class'].isin(target_classes)] other_df = df[~df['class'].isin(target_classes)] # Определяем минимальный размер среди целевых классов class_counts = target_df['class'].value_counts() min_count = class_counts.min() # Балансируем каждый целевой класс balanced_parts = [] for cls in target_classes: cls_df = target_df[target_df['class'] == cls] if len(cls_df) > min_count: cls_df = cls_df.sample(n=min_count, random_state=random_state) balanced_parts.append(cls_df) balanced_target = pd.concat(balanced_parts, ignore_index=True) # Объединяем с остальными данными (none и др.) balanced_df = pd.concat([balanced_target, other_df], ignore_index=True) return balanced_df def train(csv_path, model_path, model_type='mlp', test_size=0.2, random_state=42, balance=False, target_classes=None): # Загрузка данных df = pd.read_csv(csv_path) print(f"Total samples: {len(df)}") # Балансировка if balance: if target_classes is None: # По умолчанию балансируем все классы кроме 'none' (если есть) all_classes = df['class'].unique() target_classes = [c for c in all_classes if c != 'none'] if not target_classes: raise ValueError("No target classes found (none is the only class).") print(f"Balancing target classes: {target_classes}") df = balance_data(df, target_classes, random_state) print(f"After balancing: {len(df)} samples") print(df['class'].value_counts()) # Разделение на признаки и метки X = df.iloc[:, 1:].values.astype(np.float32) y_labels = df.iloc[:, 0].values class_names = sorted(df['class'].unique()) label_to_id = {label: i for i, label in enumerate(class_names)} y = np.array([label_to_id[label] for label in y_labels]) print(f"Classes: {class_names}") print(f"Feature count: {X.shape[1]}") # Стратифицированное разбиение X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=test_size, random_state=random_state, stratify=y ) print(f"Train size: {len(X_train)}, Test size: {len(X_test)}") # Выбор модели if model_type == 'linear': #попробовать добавить логистическую регрессию с полиномиальными признаками, можно чисто для сравнения model = LogisticRegression(max_iter=1000, random_state=random_state) elif model_type == 'mlp': model = MLPClassifier(hidden_layer_sizes=(64, 32), activation='relu', solver='adam', max_iter=500, random_state=random_state, early_stopping=True, validation_fraction=0.2) elif model_type == 'rf': model = RandomForestClassifier(n_estimators=50, max_depth=10, random_state=random_state) else: raise ValueError("model_type должен быть linear, mlp или rf") # Обучение model.fit(X_train, y_train) # Предсказание на тесте y_pred = model.predict(X_test) # Метрики accuracy = accuracy_score(y_test, y_pred) report = classification_report(y_test, y_pred, target_names=class_names, output_dict=True) conf_matrix = confusion_matrix(y_test, y_pred).tolist() print(f"\nAccuracy: {accuracy:.4f}") print("\nClassification Report:") for cls in class_names: print(f"{cls}: precision={report[cls]['precision']:.3f}, recall={report[cls]['recall']:.3f}, f1={report[cls]['f1-score']:.3f}") print("\nConfusion Matrix:") print(conf_matrix) # Сохранение модели и отчёта joblib.dump({'model': model, 'class_names': class_names}, model_path) report_data = { 'model_type': model_type, 'accuracy': accuracy, 'classification_report': report, 'confusion_matrix': conf_matrix, 'train_samples': len(X_train), 'test_samples': len(X_test), 'classes': class_names, 'balance': balance, 'target_classes': target_classes if balance else None, } with open(model_path.replace('.pkl', '_report.json'), 'w') as f: json.dump(report_data, f, indent=2) print(f"\nModel saved to {model_path}") print(f"Report saved to {model_path.replace('.pkl', '_report.json')}") if __name__ == '__main__': parser = argparse.ArgumentParser(description='Обучение модели для распознавания жестов') parser.add_argument('--csv', required=True, help='Путь к CSV-файлу с данными') parser.add_argument('--model', required=True, help='Путь для сохранения модели (.pkl)') parser.add_argument('--type', default='mlp', choices=['linear', 'mlp', 'rf'], help='Тип модели') parser.add_argument('--test_size', type=float, default=0.2, help='Доля тестовой выборки') parser.add_argument('--random_state', type=int, default=42, help='Seed для воспроизводимости') parser.add_argument('--balance', action='store_true', help='Балансировать классы (кроме none)') parser.add_argument('--target_classes', type=str, default=None, help='Список целевых классов для балансировки через запятую (по умолчанию все, кроме none)') args = parser.parse_args() target_classes = args.target_classes.split(',') if args.target_classes else None train(args.csv, args.model, args.type, args.test_size, args.random_state, args.balance, target_classes)