create_new_model.py

"""
create_new_model.py - Create a new working model from scratch
"""

import tensorflow as tf
from tensorflow.keras import layers, Model
import numpy as np
import os

print("=" * 60)
print("CREATING NEW WORKING MODEL")
print("=" * 60)

def create_vit_extractor():
    """Create a Vision Transformer feature extractor."""
    print("Building Vision Transformer extractor...")
    
    # Input layer
    inputs = layers.Input(shape=(224, 224, 3), name='input_image')
    
    # Preprocessing
    x = layers.Rescaling(1./255)(inputs)
    
    # Patch embeddings (simulated with Conv2D)
    x = layers.Conv2D(
        filters=768,
        kernel_size=16,
        strides=16,
        padding='valid',
        name='patch_embeddings'
    )(x)
    
    # Reshape to (batch, num_patches, embed_dim)
    batch_size = tf.shape(x)[0]
    x = tf.reshape(x, [batch_size, -1, 768])
    
    # Add class token
    cls_token = tf.Variable(
        tf.random.normal((1, 1, 768)),
        trainable=True,
        name='cls_token'
    )
    cls_token = tf.tile(cls_token, [batch_size, 1, 1])
    x = tf.concat([cls_token, x], axis=1)
    
    # Position embeddings
    pos_emb = layers.Dense(768, name='position_embeddings')(tf.zeros((1, 197, 768)))
    x = x + pos_emb
    
    # Transformer blocks
    for i in range(6):
        # Layer norm 1
        x_norm = layers.LayerNormalization(epsilon=1e-6, name=f'transformer_{i}_ln1')(x)
        
        # Multi-head attention
        attn = layers.MultiHeadAttention(
            num_heads=12,
            key_dim=64,
            name=f'transformer_{i}_attn'
        )
        attn_output = attn(x_norm, x_norm)
        
        # First residual
        x = x + attn_output
        
        # Layer norm 2
        x_norm = layers.LayerNormalization(epsilon=1e-6, name=f'transformer_{i}_ln2')(x)
        
        # MLP
        mlp = tf.keras.Sequential([
            layers.Dense(3072, activation='gelu', name=f'transformer_{i}_mlp1'),
            layers.Dropout(0.1),
            layers.Dense(768, name=f'transformer_{i}_mlp2'),
            layers.Dropout(0.1)
        ], name=f'transformer_{i}_mlp')
        
        mlp_output = mlp(x_norm)
        
        # Second residual
        x = x + mlp_output
    
    # Extract class token
    x = x[:, 0, :]  # Shape: (batch, 768)
    
    # Project to feature space
    x = layers.Dense(512, activation='relu', name='proj1')(x)
    x = layers.Dropout(0.3, name='dropout1')(x)
    x = layers.Dense(256, activation='relu', name='proj2')(x)
    outputs = layers.Dense(512, activation='relu', name='output_features')(x)
    
    # Create model
    model = Model(inputs=inputs, outputs=outputs, name='hybrid_feature_extractor')
    
    # Compile (not needed for inference but helps)
    model.compile(optimizer='adam', loss='mse')
    
    print(f"✓ Vision Transformer created")
    return model

def create_cnn_extractor():
    """Create a CNN-based feature extractor (faster)."""
    print("Building CNN extractor...")
    
    inputs = layers.Input(shape=(224, 224, 3), name='input_image')
    
    # Use EfficientNetV2B0 for good performance
    base_model = tf.keras.applications.EfficientNetV2B0(
        include_top=False,
        weights='imagenet',
        input_tensor=inputs,
        pooling='avg'
    )
    base_model.trainable = False
    
    # Get features
    features = base_model.output
    
    # Add custom layers
    x = layers.Dense(512, activation='relu', name='fc1')(features)
    x = layers.Dropout(0.3, name='dropout1')(x)
    x = layers.Dense(256, activation='relu', name='fc2')(x)
    x = layers.Dropout(0.2, name='dropout2')(x)
    outputs = layers.Dense(512, activation='relu', name='output_features')(x)
    
    model = Model(inputs=inputs, outputs=outputs, name='cnn_feature_extractor')
    
    print(f"✓ CNN extractor created")
    return model

def save_model_multiple_formats(model, base_name):
    """Save model in multiple formats for compatibility."""
    print(f"\nSaving model in multiple formats...")
    
    # 1. Save as .keras (recommended)
    keras_path = f"{base_name}.keras"
    model.save(keras_path)
    print(f"✓ Saved as {keras_path}")
    
    # 2. Save as SavedModel
    sm_path = f"{base_name}_savedmodel"
    model.save(sm_path, save_format='tf')
    print(f"✓ Saved as SavedModel: {sm_path}")
    
    # 3. Save as H5 (for backward compatibility)
    h5_path = f"{base_name}.h5"
    model.save(h5_path, save_format='h5')
    print(f"✓ Saved as H5: {h5_path}")
    
    return keras_path, sm_path, h5_path

def create_dummy_data(model):
    """Create dummy training data to initialize model."""
    print("\nCreating dummy training data...")
    
    # Generate random images
    dummy_images = np.random.randn(10, 224, 224, 3).astype('float32')
    dummy_features = np.random.randn(10, model.output_shape[1]).astype('float32')
    
    # Train for 1 epoch (just to initialize weights properly)
    model.fit(
        dummy_images,
        dummy_features,
        epochs=1,
        verbose=0
    )
    
    print("✓ Model initialized with dummy data")

if __name__ == "__main__":
    print("Select model type:")
    print("1. Vision Transformer (ViT) - More accurate")
    print("2. CNN (EfficientNet) - Faster inference")
    print("3. Hybrid (CNN + Transformer) - Best of both")
    
    choice = input("\nEnter choice (1, 2, or 3): ").strip()
    
    if choice == '1':
        model = create_vit_extractor()
        model_name = "vit_extractor"
    elif choice == '2':
        model = create_cnn_extractor()
        model_name = "cnn_extractor"
    else:
        # Create hybrid model
        print("Building hybrid model...")
        model = create_cnn_extractor()
        model_name = "hybrid_extractor"
    
    # Show model summary
    print(f"\nModel Summary:")
    model.summary()
    
    # Initialize with dummy data
    create_dummy_data(model)
    
    # Save in multiple formats
    keras_path, sm_path, h5_path = save_model_multiple_formats(model, model_name)
    
    print("\n" + "=" * 60)
    print("MODEL CREATION COMPLETE")
    print("=" * 60)
    print(f"✅ Model: {model.name}")
    print(f"✅ Input shape: {model.input_shape}")
    print(f"✅ Output shape: {model.output_shape}")
    print(f"✅ Formats available:")
    print(f"   - {keras_path} (recommended)")
    print(f"   - {sm_path}/ (SavedModel)")
    print(f"   - {h5_path} (H5 for compatibility)")
    print("\nNow update your app.py to use one of these files!")
    print("=" * 60)