train_model.py

"""
TinyML Training Pipeline for Slip Detection
Uses Decision Tree classifier optimized for Arduino Uno deployment
"""

import numpy as np
import pandas as pd
from sklearn.tree import DecisionTreeClassifier, export_text
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import joblib
import json
import os

class SlipDetectionTrainer:
    def __init__(self, max_depth=5, min_samples_split=10, min_samples_leaf=5):
        """
        Initialize trainer with parameters optimized for Arduino Uno
        
        Args:
            max_depth: Maximum depth of tree (lower = smaller model)
            min_samples_split: Minimum samples to split a node
            min_samples_leaf: Minimum samples in a leaf node
        """
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.min_samples_leaf = min_samples_leaf
        self.model = None
        self.feature_names = None
        
    def load_data(self, data_path):
        """
        Load training data from CSV file
        
        Expected CSV format:
        - Columns: feature1, feature2, ..., featureN, label
        - Label: 0 = no slip, 1 = slip
        """
        df = pd.read_csv(data_path)
        
        # Separate features and labels
        self.feature_names = [col for col in df.columns if col != 'label']
        X = df[self.feature_names].values
        y = df['label'].values
        
        # Convert labels to integers (handle float labels like 0.0, 1.0)
        y = y.astype(int)
        
        print(f"Loaded {len(X)} samples with {len(self.feature_names)} features")
        print(f"Features: {self.feature_names}")
        print(f"Class distribution: {np.bincount(y)}")
        
        return X, y
    
    def train(self, X, y, test_size=0.2, random_state=42):
        """
        Train decision tree model
        """
        # Split data
        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"\nTraining set: {len(X_train)} samples")
        print(f"Test set: {len(X_test)} samples")
        
        # Create and train model
        self.model = DecisionTreeClassifier(
            max_depth=self.max_depth,
            min_samples_split=self.min_samples_split,
            min_samples_leaf=self.min_samples_leaf,
            random_state=random_state,
            criterion='gini'
        )
        
        print("\nTraining model...")
        self.model.fit(X_train, y_train)
        
        # Evaluate
        train_acc = accuracy_score(y_train, self.model.predict(X_train))
        test_acc = accuracy_score(y_test, self.model.predict(X_test))
        
        print(f"\nTraining Accuracy: {train_acc:.4f}")
        print(f"Test Accuracy: {test_acc:.4f}")
        
        print("\nClassification Report:")
        print(classification_report(y_test, self.model.predict(X_test)))
        
        print("\nConfusion Matrix:")
        print(confusion_matrix(y_test, self.model.predict(X_test)))
        
        # Print tree structure
        print("\nDecision Tree Structure:")
        print(export_text(self.model, feature_names=self.feature_names, max_depth=10))
        
        return X_test, y_test
    
    def save_model(self, model_path='model.pkl'):
        """Save trained model"""
        if self.model is None:
            raise ValueError("Model not trained yet")
        joblib.dump(self.model, model_path)
        print(f"\nModel saved to {model_path}")
    
    def export_for_arduino(self, output_path='arduino_model.h'):
        """
        Export decision tree as Arduino C++ code
        """
        if self.model is None:
            raise ValueError("Model not trained yet")
        
        tree = self.model.tree_
        
        # Generate Arduino header file
        code = self._generate_arduino_code(tree)
        
        with open(output_path, 'w') as f:
            f.write(code)
        
        print(f"\nArduino model exported to {output_path}")
        
        # Also save metadata
        metadata = {
            'n_features': int(len(self.feature_names)),
            'feature_names': self.feature_names,
            'n_nodes': int(tree.node_count),
            'max_depth': int(self.max_depth),
            'n_classes': int(tree.n_classes[0])
        }
        
        with open('model_metadata.json', 'w') as f:
            json.dump(metadata, f, indent=2)
        
        print(f"Model metadata saved to model_metadata.json")
    
    def _generate_arduino_code(self, tree):
        """Generate Arduino C++ code for decision tree inference"""
        
        # Get tree structure
        children_left = tree.children_left.tolist()
        children_right = tree.children_right.tolist()
        feature = tree.feature.tolist()
        threshold = tree.threshold.tolist()
        value = tree.value.tolist()
        
        # Count nodes
        n_nodes = tree.node_count
        
        code = f"""/*
 * Auto-generated Decision Tree Model for Slip Detection
 * Generated from scikit-learn DecisionTreeClassifier
 * 
 * Model Statistics:
 * - Nodes: {n_nodes}
 * - Features: {len(self.feature_names)}
 * - Max Depth: {self.max_depth}
 */

#ifndef SLIP_DETECTION_MODEL_H
#define SLIP_DETECTION_MODEL_H

#include <Arduino.h>

// Model configuration
#define N_NODES {n_nodes}
#define N_FEATURES {len(self.feature_names)}

// Decision tree structure
static const int children_left[N_NODES] = {{ {', '.join(map(str, children_left))} }};
static const int children_right[N_NODES] = {{ {', '.join(map(str, children_right))} }};
static const int feature[N_NODES] = {{ {', '.join(map(str, feature))} }};
static const float threshold[N_NODES] = {{ {', '.join([f'{t:.6f}f' for t in threshold])} }};

// Leaf node predictions (value[node][class])
static const float value[N_NODES][2] = {{
"""
        
        # Add value array
        for i, val in enumerate(value):
            # Normalize probabilities
            total = val[0][0] + val[0][1] if len(val[0]) > 1 else val[0][0]
            prob_0 = val[0][0] / total if total > 0 else 0.0
            prob_1 = val[0][1] / total if len(val[0]) > 1 else 0.0
            
            code += f"    {{ {prob_0:.6f}f, {prob_1:.6f}f }},  // Node {i}\n"
        
        code += """};

/**
 * Predict slip detection from sensor features
 * 
 * @param features Array of N_FEATURES sensor readings
 * @return 0 = no slip, 1 = slip detected
 */
int predict_slip(float features[N_FEATURES]) {
    int node = 0;
    
    while (children_left[node] != children_right[node]) {
        // Not a leaf node, traverse tree
        if (features[feature[node]] <= threshold[node]) {
            node = children_left[node];
        } else {
            node = children_right[node];
        }
    }
    
    // Leaf node reached, return prediction
    return (value[node][1] > value[node][0]) ? 1 : 0;
}

/**
 * Get prediction probability
 * 
 * @param features Array of N_FEATURES sensor readings
 * @return Probability of slip (0.0 to 1.0)
 */
float predict_slip_probability(float features[N_FEATURES]) {
    int node = 0;
    
    while (children_left[node] != children_right[node]) {
        if (features[feature[node]] <= threshold[node]) {
            node = children_left[node];
        } else {
            node = children_right[node];
        }
    }
    
    return value[node][1];
}

#endif // SLIP_DETECTION_MODEL_H
"""
        
        return code


def main():
    """Main training pipeline"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Train slip detection model')
    parser.add_argument('--data', type=str, default='data/training_data.csv',
                       help='Path to training data CSV')
    parser.add_argument('--max-depth', type=int, default=5,
                       help='Maximum tree depth')
    parser.add_argument('--min-samples-split', type=int, default=10,
                       help='Minimum samples to split')
    parser.add_argument('--min-samples-leaf', type=int, default=5,
                       help='Minimum samples in leaf')
    parser.add_argument('--output-model', type=str, default='model.pkl',
                       help='Output path for saved model')
    parser.add_argument('--output-arduino', type=str, default='arduino_model.h',
                       help='Output path for Arduino header')
    
    args = parser.parse_args()
    
    # Create trainer
    trainer = SlipDetectionTrainer(
        max_depth=args.max_depth,
        min_samples_split=args.min_samples_split,
        min_samples_leaf=args.min_samples_leaf
    )
    
    # Load and train
    if not os.path.exists(args.data):
        print(f"Error: Data file {args.data} not found!")
        print("Please create training data or use generate_sample_data.py")
        return
    
    X, y = trainer.load_data(args.data)
    trainer.train(X, y)
    
    # Save model
    trainer.save_model(args.output_model)
    
    # Export for Arduino
    trainer.export_for_arduino(args.output_arduino)
    
    print("\n✓ Training pipeline completed successfully!")
    print(f"✓ Model saved: {args.output_model}")
    print(f"✓ Arduino code: {args.output_arduino}")


if __name__ == '__main__':
    main()