src/models/mobilenet_model.py

"""
MobileNetV2 transfer learning model for emotion recognition.
"""
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import (
    Dense, Dropout, GlobalAveragePooling2D,
    BatchNormalization, Input, Lambda
)
from tensorflow.keras.applications import MobileNetV2

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from src.config import IMAGE_SIZE_TRANSFER, NUM_CLASSES, NUM_CHANNELS_RGB


def build_mobilenet_model(
    input_shape: tuple = (*IMAGE_SIZE_TRANSFER, NUM_CHANNELS_RGB),
    num_classes: int = NUM_CLASSES,
    trainable_layers: int = 30,
    dropout_rate: float = 0.5
) -> Model:
    """
    Build MobileNetV2 transfer learning model for emotion recognition.
    
    Args:
        input_shape: Input image shape (height, width, channels)
        num_classes: Number of emotion classes
        trainable_layers: Number of top layers to make trainable
        dropout_rate: Dropout rate for dense layers
        
    Returns:
        Keras model
    """
    # Load pre-trained MobileNetV2
    base_model = MobileNetV2(
        weights='imagenet',
        include_top=False,
        input_shape=input_shape
    )
    
    # Freeze base layers
    for layer in base_model.layers[:-trainable_layers]:
        layer.trainable = False
    
    # Make top layers trainable
    for layer in base_model.layers[-trainable_layers:]:
        layer.trainable = True
    
    # Build the model
    inputs = Input(shape=input_shape)
    
    # Preprocess input for MobileNetV2 using Rescaling layer
    # MobileNetV2 expects inputs in [-1, 1] range
    x = tf.keras.layers.Rescaling(scale=1./127.5, offset=-1.0)(inputs)
    
    # Pass through base model
    x = base_model(x, training=True)
    
    # Classification head
    x = GlobalAveragePooling2D()(x)
    x = Dense(512, activation='relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(dropout_rate)(x)
    x = Dense(256, activation='relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(dropout_rate)(x)
    outputs = Dense(num_classes, activation='softmax')(x)
    
    model = Model(inputs=inputs, outputs=outputs, name='mobilenet_emotion')
    
    return model


def build_mobilenet_from_grayscale(
    input_shape: tuple = (*IMAGE_SIZE_TRANSFER, 1),
    num_classes: int = NUM_CLASSES,
    trainable_layers: int = 30,
    dropout_rate: float = 0.5
) -> Model:
    """
    Build MobileNetV2 model that accepts grayscale input.
    Converts grayscale to RGB internally.
    
    Args:
        input_shape: Input shape for grayscale images
        num_classes: Number of emotion classes
        trainable_layers: Number of top layers to make trainable
        dropout_rate: Dropout rate
        
    Returns:
        Keras model
    """
    # Load pre-trained MobileNetV2
    base_model = MobileNetV2(
        weights='imagenet',
        include_top=False,
        input_shape=(*IMAGE_SIZE_TRANSFER, 3)
    )
    
    # Freeze base layers
    for layer in base_model.layers[:-trainable_layers]:
        layer.trainable = False
    
    # Input for grayscale image
    inputs = Input(shape=input_shape)
    
    # Convert grayscale to RGB by repeating channels
    x = tf.keras.layers.Concatenate()([inputs, inputs, inputs])
    
    # Preprocess for MobileNetV2 using Rescaling layer
    x = tf.keras.layers.Rescaling(scale=1./127.5, offset=-1.0)(x)
    
    # Base model
    x = base_model(x, training=True)
    
    # Classification head
    x = GlobalAveragePooling2D()(x)
    x = Dense(512, activation='relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(dropout_rate)(x)
    x = Dense(256, activation='relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(dropout_rate)(x)
    outputs = Dense(num_classes, activation='softmax')(x)
    
    model = Model(inputs=inputs, outputs=outputs, name='mobilenet_emotion_grayscale')
    
    return model


def freeze_base_model(model: Model) -> Model:
    """
    Freeze all layers in the base MobileNetV2 model.
    Useful for initial training with frozen weights.
    
    Args:
        model: MobileNet emotion model
        
    Returns:
        Model with frozen base
    """
    for layer in model.layers:
        if 'mobilenet' in layer.name.lower():
            layer.trainable = False
    return model


def unfreeze_top_layers(model: Model, num_layers: int = 30) -> Model:
    """
    Unfreeze top layers of the base model for fine-tuning.
    
    Args:
        model: MobileNet emotion model
        num_layers: Number of top layers to unfreeze
        
    Returns:
        Model with partially unfrozen base
    """
    for layer in model.layers:
        if 'mobilenet' in layer.name.lower():
            # Get base model and unfreeze top layers
            for base_layer in layer.layers[-num_layers:]:
                base_layer.trainable = True
    return model


def get_model_config() -> dict:
    """
    Get the default model configuration.
    
    Returns:
        Dictionary with model configuration
    """
    return {
        "name": "MobileNetV2",
        "input_shape": (*IMAGE_SIZE_TRANSFER, NUM_CHANNELS_RGB),
        "num_classes": NUM_CLASSES,
        "expected_accuracy": "65-72%",
        "training_time": "~45 minutes (GPU)",
        "parameters": "~3.5M",
        "base_model": "MobileNetV2 (ImageNet)"
    }


if __name__ == "__main__":
    # Build and display model summary
    print("Building MobileNetV2 model...")
    model = build_mobilenet_model()
    
    # Count trainable parameters
    trainable = sum([tf.keras.backend.count_params(w) for w in model.trainable_weights])
    non_trainable = sum([tf.keras.backend.count_params(w) for w in model.non_trainable_weights])
    
    print(f"\nTotal parameters: {trainable + non_trainable:,}")
    print(f"Trainable parameters: {trainable:,}")
    print(f"Non-trainable parameters: {non_trainable:,}")
    
    print("\nModel configuration:")
    config = get_model_config()
    for key, value in config.items():
        print(f"  {key}: {value}")