train_model.py

#!/usr/bin/env python3
"""
SkinMate ML Model Training Script
Advanced skin type classification using TensorFlow/Keras
"""

import os
import sys
import time
import json
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Dict, Tuple, List
import cv2
from PIL import Image

# ML Libraries
try:
    import tensorflow as tf
    from tensorflow import keras
    from tensorflow.keras import layers, optimizers, callbacks
    from tensorflow.keras.preprocessing.image import ImageDataGenerator
    from tensorflow.keras.applications import EfficientNetB0, MobileNetV2, ResNet50V2
    from sklearn.metrics import classification_report, confusion_matrix
    from sklearn.utils.class_weight import compute_class_weight
    import seaborn as sns
    print("✅ All ML libraries available")
except ImportError as e:
    print(f"⚠️ Installing required ML libraries: {e}")
    os.system("pip install tensorflow scikit-learn matplotlib seaborn")
    import tensorflow as tf
    from tensorflow import keras
    from tensorflow.keras import layers, optimizers, callbacks
    from tensorflow.keras.preprocessing.image import ImageDataGenerator
    from tensorflow.keras.applications import EfficientNetB0, MobileNetV2, ResNet50V2
    from sklearn.metrics import classification_report, confusion_matrix
    from sklearn.utils.class_weight import compute_class_weight
    import seaborn as sns

class SkinMateModelTrainer:
    """Advanced skin type classification model trainer"""
    
    def __init__(self, dataset_dir: Path, model_output_dir: Path):
        self.dataset_dir = dataset_dir
        self.model_output_dir = model_output_dir
        self.model_output_dir.mkdir(parents=True, exist_ok=True)
        
        # Model configuration
        self.IMG_SIZE = 224
        self.BATCH_SIZE = 32
        self.NUM_CLASSES = 3  # normal, oily, dry (combination and sensitive have no data yet)
        self.CLASS_NAMES = ['dry', 'normal', 'oily']
        
        # Training configuration
        self.EPOCHS = 50
        self.LEARNING_RATE = 0.001
        self.PATIENCE = 10  # Early stopping patience
        
        # Initialize TensorFlow
        self.setup_tensorflow()
        
    def setup_tensorflow(self):
        """Setup TensorFlow configuration"""
        # GPU configuration
        gpus = tf.config.experimental.list_physical_devices('GPU')
        if gpus:
            try:
                for gpu in gpus:
                    tf.config.experimental.set_memory_growth(gpu, True)
                print(f"✅ GPU available: {len(gpus)} GPU(s)")
            except RuntimeError as e:
                print(f"⚠️ GPU setup error: {e}")
        else:
            print("ℹ️ Running on CPU")
        
        # Mixed precision for better performance
        try:
            policy = tf.keras.mixed_precision.Policy('mixed_float16')
            tf.keras.mixed_precision.set_global_policy(policy)
            print("✅ Mixed precision enabled")
        except:
            print("ℹ️ Mixed precision not available")
    
    def create_data_generators(self) -> Tuple[tf.keras.utils.Sequence, tf.keras.utils.Sequence, tf.keras.utils.Sequence]:
        """Create data generators with augmentation"""
        
        # Data augmentation for training
        train_datagen = ImageDataGenerator(
            rescale=1./255,
            rotation_range=20,
            width_shift_range=0.1,
            height_shift_range=0.1,
            shear_range=0.1,
            zoom_range=0.1,
            horizontal_flip=True,
            brightness_range=[0.8, 1.2],
            fill_mode='nearest'
        )
        
        # No augmentation for validation and test
        val_test_datagen = ImageDataGenerator(rescale=1./255)
        
        # Create generators
        train_generator = train_datagen.flow_from_directory(
            self.dataset_dir / 'train',
            target_size=(self.IMG_SIZE, self.IMG_SIZE),
            batch_size=self.BATCH_SIZE,
            class_mode='categorical',
            classes=self.CLASS_NAMES,
            shuffle=True,
            seed=42
        )
        
        validation_generator = val_test_datagen.flow_from_directory(
            self.dataset_dir / 'validation',
            target_size=(self.IMG_SIZE, self.IMG_SIZE),
            batch_size=self.BATCH_SIZE,
            class_mode='categorical',
            classes=self.CLASS_NAMES,
            shuffle=False,
            seed=42
        )
        
        test_generator = val_test_datagen.flow_from_directory(
            self.dataset_dir / 'test',
            target_size=(self.IMG_SIZE, self.IMG_SIZE),
            batch_size=self.BATCH_SIZE,
            class_mode='categorical',
            classes=self.CLASS_NAMES,
            shuffle=False,
            seed=42
        )
        
        print(f"✅ Data generators created:")
        print(f"   Training samples: {train_generator.samples}")
        print(f"   Validation samples: {validation_generator.samples}")
        print(f"   Test samples: {test_generator.samples}")
        print(f"   Classes: {train_generator.class_indices}")
        
        return train_generator, validation_generator, test_generator
    
    def create_advanced_model(self, model_type: str = "efficientnet") -> keras.Model:
        """Create advanced CNN model with transfer learning"""
        
        # Input layer
        inputs = keras.Input(shape=(self.IMG_SIZE, self.IMG_SIZE, 3))
        
        # Data augmentation layer (applied during training)
        x = layers.experimental.preprocessing.RandomFlip("horizontal")(inputs)
        x = layers.experimental.preprocessing.RandomRotation(0.1)(x)
        x = layers.experimental.preprocessing.RandomZoom(0.1)(x)
        
        # Base model selection
        if model_type == "efficientnet":
            base_model = EfficientNetB0(
                input_shape=(self.IMG_SIZE, self.IMG_SIZE, 3),
                include_top=False,
                weights='imagenet'
            )
        elif model_type == "mobilenet":
            base_model = MobileNetV2(
                input_shape=(self.IMG_SIZE, self.IMG_SIZE, 3),
                include_top=False,
                weights='imagenet'
            )
        else:  # resnet
            base_model = ResNet50V2(
                input_shape=(self.IMG_SIZE, self.IMG_SIZE, 3),
                include_top=False,
                weights='imagenet'
            )
        
        # Freeze base model initially
        base_model.trainable = False
        
        # Apply base model
        x = base_model(x, training=False)
        
        # Add custom head
        x = layers.GlobalAveragePooling2D()(x)
        x = layers.Dropout(0.3)(x)
        x = layers.Dense(512, activation='relu')(x)
        x = layers.BatchNormalization()(x)
        x = layers.Dropout(0.2)(x)
        x = layers.Dense(256, activation='relu')(x)
        x = layers.BatchNormalization()(x)
        x = layers.Dropout(0.1)(x)
        
        # Output layer
        outputs = layers.Dense(self.NUM_CLASSES, activation='softmax', dtype='float32')(x)
        
        model = keras.Model(inputs, outputs)
        
        print(f"✅ {model_type.title()} model created")
        print(f"   Total parameters: {model.count_params():,}")
        print(f"   Trainable parameters: {sum([keras.backend.count_params(w) for w in model.trainable_weights]):,}")
        
        return model, base_model
    
    def calculate_class_weights(self, train_generator) -> Dict[int, float]:
        """Calculate class weights for imbalanced dataset"""
        
        # Get class distribution
        class_counts = {}
        for class_name in self.CLASS_NAMES:
            class_dir = self.dataset_dir / 'train' / class_name
            class_counts[class_name] = len(list(class_dir.glob('*.jpg')))
        
        print(f"📊 Class distribution: {class_counts}")
        
        # Calculate weights
        total_samples = sum(class_counts.values())
        class_weights = {}
        
        for i, class_name in enumerate(self.CLASS_NAMES):
            class_weights[i] = total_samples / (len(self.CLASS_NAMES) * class_counts[class_name])
        
        print(f"⚖️ Class weights: {class_weights}")
        return class_weights
    
    def create_callbacks(self) -> List[keras.callbacks.Callback]:
        """Create training callbacks"""
        
        callbacks_list = [
            # Early stopping
            keras.callbacks.EarlyStopping(
                monitor='val_accuracy',
                patience=self.PATIENCE,
                restore_best_weights=True,
                verbose=1
            ),
            
            # Learning rate reduction
            keras.callbacks.ReduceLROnPlateau(
                monitor='val_loss',
                factor=0.2,
                patience=5,
                min_lr=1e-7,
                verbose=1
            ),
            
            # Model checkpoint
            keras.callbacks.ModelCheckpoint(
                filepath=str(self.model_output_dir / 'best_model.h5'),
                monitor='val_accuracy',
                save_best_only=True,
                save_weights_only=False,
                verbose=1
            ),
            
            # TensorBoard
            keras.callbacks.TensorBoard(
                log_dir=str(self.model_output_dir / 'logs'),
                histogram_freq=1,
                write_graph=True,
                write_images=True
            )
        ]
        
        return callbacks_list
    
    def train_model(self, model: keras.Model, base_model: keras.Model, 
                   train_gen, val_gen, class_weights: Dict) -> keras.callbacks.History:
        """Train the model with fine-tuning"""
        
        print("\\n🚀 Starting model training...")
        
        # Compile model for initial training
        model.compile(
            optimizer=optimizers.Adam(learning_rate=self.LEARNING_RATE),
            loss='categorical_crossentropy',
            metrics=['accuracy', 'top_2_accuracy']
        )
        
        # Phase 1: Train only the head
        print("\\n📚 Phase 1: Training classifier head...")
        
        history1 = model.fit(
            train_gen,
            epochs=15,
            validation_data=val_gen,
            class_weight=class_weights,
            callbacks=self.create_callbacks(),
            verbose=1
        )
        
        # Phase 2: Fine-tune the entire model
        print("\\n🔧 Phase 2: Fine-tuning entire model...")
        
        # Unfreeze base model
        base_model.trainable = True
        
        # Use lower learning rate for fine-tuning
        model.compile(
            optimizer=optimizers.Adam(learning_rate=self.LEARNING_RATE/10),
            loss='categorical_crossentropy',
            metrics=['accuracy', 'top_2_accuracy']
        )
        
        history2 = model.fit(
            train_gen,
            epochs=self.EPOCHS - 15,
            validation_data=val_gen,
            class_weight=class_weights,
            callbacks=self.create_callbacks(),
            verbose=1,
            initial_epoch=15
        )
        
        # Combine histories
        combined_history = {}
        for key in history1.history.keys():
            combined_history[key] = history1.history[key] + history2.history[key]
        
        return combined_history
    
    def evaluate_model(self, model: keras.Model, test_gen) -> Dict:
        """Comprehensive model evaluation"""
        
        print("\\n📊 Evaluating model...")
        
        # Test set evaluation
        test_loss, test_accuracy, test_top2 = model.evaluate(test_gen, verbose=1)
        
        # Predictions
        test_gen.reset()
        y_pred = model.predict(test_gen, verbose=1)
        y_pred_classes = np.argmax(y_pred, axis=1)
        y_true = test_gen.classes
        
        # Classification report
        report = classification_report(
            y_true, y_pred_classes, 
            target_names=self.CLASS_NAMES,
            output_dict=True
        )
        
        # Confusion matrix
        cm = confusion_matrix(y_true, y_pred_classes)
        
        results = {
            'test_loss': float(test_loss),
            'test_accuracy': float(test_accuracy),
            'test_top2_accuracy': float(test_top2),
            'classification_report': report,
            'confusion_matrix': cm.tolist(),
            'class_names': self.CLASS_NAMES
        }
        
        # Print results
        print(f"\\n✅ Model Evaluation Results:")
        print(f"   Test Accuracy: {test_accuracy:.4f}")
        print(f"   Test Top-2 Accuracy: {test_top2:.4f}")
        print(f"   Test Loss: {test_loss:.4f}")
        
        print(f"\\n📋 Classification Report:")
        for class_name in self.CLASS_NAMES:
            metrics = report[class_name]
            print(f"   {class_name.title()}:")
            print(f"      Precision: {metrics['precision']:.3f}")
            print(f"      Recall: {metrics['recall']:.3f}")
            print(f"      F1-Score: {metrics['f1-score']:.3f}")
        
        return results
    
    def plot_training_history(self, history: Dict):
        """Plot training history"""
        
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        
        # Accuracy
        axes[0, 0].plot(history['accuracy'], label='Training Accuracy')
        axes[0, 0].plot(history['val_accuracy'], label='Validation Accuracy')
        axes[0, 0].set_title('Model Accuracy')
        axes[0, 0].set_xlabel('Epoch')
        axes[0, 0].set_ylabel('Accuracy')
        axes[0, 0].legend()
        axes[0, 0].grid(True)
        
        # Loss
        axes[0, 1].plot(history['loss'], label='Training Loss')
        axes[0, 1].plot(history['val_loss'], label='Validation Loss')
        axes[0, 1].set_title('Model Loss')
        axes[0, 1].set_xlabel('Epoch')
        axes[0, 1].set_ylabel('Loss')
        axes[0, 1].legend()
        axes[0, 1].grid(True)
        
        # Top-2 Accuracy
        axes[1, 0].plot(history['top_2_accuracy'], label='Training Top-2 Accuracy')
        axes[1, 0].plot(history['val_top_2_accuracy'], label='Validation Top-2 Accuracy')
        axes[1, 0].set_title('Model Top-2 Accuracy')
        axes[1, 0].set_xlabel('Epoch')
        axes[1, 0].set_ylabel('Top-2 Accuracy')
        axes[1, 0].legend()
        axes[1, 0].grid(True)
        
        # Learning rate
        if 'lr' in history:
            axes[1, 1].plot(history['lr'], label='Learning Rate')
            axes[1, 1].set_title('Learning Rate Schedule')
            axes[1, 1].set_xlabel('Epoch')
            axes[1, 1].set_ylabel('Learning Rate')
            axes[1, 1].set_yscale('log')
            axes[1, 1].legend()
            axes[1, 1].grid(True)
        else:
            axes[1, 1].axis('off')
        
        plt.tight_layout()
        plt.savefig(self.model_output_dir / 'training_history.png', dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"✅ Training history plot saved: {self.model_output_dir / 'training_history.png'}")
    
    def plot_confusion_matrix(self, cm: np.ndarray):
        """Plot confusion matrix"""
        
        plt.figure(figsize=(8, 6))
        sns.heatmap(
            cm, 
            annot=True, 
            fmt='d', 
            cmap='Blues',
            xticklabels=self.CLASS_NAMES,
            yticklabels=self.CLASS_NAMES
        )
        plt.title('Confusion Matrix')
        plt.xlabel('Predicted')
        plt.ylabel('Actual')
        plt.tight_layout()
        plt.savefig(self.model_output_dir / 'confusion_matrix.png', dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"✅ Confusion matrix plot saved: {self.model_output_dir / 'confusion_matrix.png'}")
    
    def save_model_artifacts(self, model: keras.Model, history: Dict, results: Dict):
        """Save all model artifacts"""
        
        # Save model in multiple formats
        model.save(self.model_output_dir / 'skinmate_model.h5')
        model.save(self.model_output_dir / 'skinmate_model', save_format='tf')
        
        # Save model weights
        model.save_weights(self.model_output_dir / 'skinmate_weights.h5')
        
        # Save training history
        with open(self.model_output_dir / 'training_history.json', 'w') as f:
            # Convert numpy arrays to lists for JSON serialization
            history_serializable = {}
            for key, value in history.items():
                if isinstance(value, (list, np.ndarray)):
                    history_serializable[key] = [float(x) for x in value]
                else:
                    history_serializable[key] = value
            json.dump(history_serializable, f, indent=2)
        
        # Save evaluation results
        with open(self.model_output_dir / 'evaluation_results.json', 'w') as f:
            json.dump(results, f, indent=2)
        
        # Save model architecture
        with open(self.model_output_dir / 'model_architecture.json', 'w') as f:
            f.write(model.to_json())
        
        # Save model summary
        with open(self.model_output_dir / 'model_summary.txt', 'w') as f:
            model.summary(print_fn=lambda x: f.write(x + '\\n'))
        
        # Save class names
        with open(self.model_output_dir / 'class_names.json', 'w') as f:
            json.dump(self.CLASS_NAMES, f)
        
        print(f"\\n✅ All model artifacts saved to: {self.model_output_dir}")
        print(f"   - Model: skinmate_model.h5")
        print(f"   - Weights: skinmate_weights.h5")
        print(f"   - TensorFlow SavedModel: skinmate_model/")
        print(f"   - Training history: training_history.json")
        print(f"   - Evaluation results: evaluation_results.json")
        print(f"   - Model architecture: model_architecture.json")
        print(f"   - Class names: class_names.json")
    
    def train_complete_pipeline(self) -> Dict:
        """Complete training pipeline"""
        
        start_time = time.time()
        
        print("🚀 SkinMate Model Training Pipeline")
        print("=" * 60)
        
        # Create data generators
        train_gen, val_gen, test_gen = self.create_data_generators()
        
        # Calculate class weights
        class_weights = self.calculate_class_weights(train_gen)
        
        # Create model
        model, base_model = self.create_advanced_model("efficientnet")
        
        # Train model
        history = self.train_model(model, base_model, train_gen, val_gen, class_weights)
        
        # Load best model
        best_model = keras.models.load_model(self.model_output_dir / 'best_model.h5')
        
        # Evaluate model
        results = self.evaluate_model(best_model, test_gen)
        
        # Create visualizations
        self.plot_training_history(history)
        self.plot_confusion_matrix(np.array(results['confusion_matrix']))
        
        # Save artifacts
        self.save_model_artifacts(best_model, history, results)
        
        # Final summary
        training_time = time.time() - start_time
        print(f"\\n🎉 TRAINING COMPLETED!")
        print("=" * 60)
        print(f"⏱️ Total training time: {training_time/3600:.2f} hours")
        print(f"🎯 Final test accuracy: {results['test_accuracy']:.4f}")
        print(f"🏆 Model ready for production!")
        
        return {
            'model_path': str(self.model_output_dir / 'skinmate_model.h5'),
            'training_time_hours': training_time / 3600,
            'final_accuracy': results['test_accuracy'],
            'results': results
        }

def main():
    """Main training execution"""
    
    # Configuration
    base_dir = Path(".")
    dataset_dir = base_dir / "datasets" / "final-combined-dataset"
    model_output_dir = base_dir / "models" / f"skinmate_model_{int(time.time())}"
    
    # Check if dataset exists
    if not dataset_dir.exists():
        print(f"❌ Dataset not found: {dataset_dir}")
        print("💡 Run prepare_dataset.py first to create the dataset")
        return
    
    # Check if we have all required classes
    required_dirs = ['train/normal', 'train/oily', 'train/dry']
    for req_dir in required_dirs:
        if not (dataset_dir / req_dir).exists():
            print(f"❌ Required directory not found: {dataset_dir / req_dir}")
            return
    
    print(f"✅ Dataset found: {dataset_dir}")
    print(f"📁 Model output: {model_output_dir}")
    
    # Initialize trainer
    trainer = SkinMateModelTrainer(dataset_dir, model_output_dir)
    
    try:
        # Run complete training pipeline
        results = trainer.train_complete_pipeline()
        
        print(f"\\n🚀 Training completed successfully!")
        print(f"📄 Model saved at: {results['model_path']}")
        
    except Exception as e:
        print(f"\\n❌ Training failed: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()