train.py

import os
import argparse
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.applications import (VGG16, VGG19, ResNet50, ResNet101, InceptionV3,
                                           DenseNet121, DenseNet201, MobileNetV2, Xception, InceptionResNetV2,
                                           NASNetLarge, NASNetMobile, EfficientNetB0, EfficientNetB7)
from tensorflow.keras.layers import Dense, Flatten, Dropout, BatchNormalization
from tensorflow.keras.optimizers import Adam, SGD
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
import numpy as np

def load_and_preprocess_image(filename, label, image_size):
    # Load image
    image = tf.io.read_file(filename)
    image = tf.image.decode_image(image, channels=3)
    
    # Ensure the image tensor has shape
    if not tf.executing_eagerly():
        image.set_shape([None, None, 3])

    # Resize image to the specified size
    image = tf.image.resize(image, [image_size[0], image_size[1]])  # Use height and width from the tuple
    
    # Normalize image to [0, 1]
    image = image / 255.0
    image.set_shape([image_size[0], image_size[1], 3])

    return image, label

def create_dataset(data_dir, labels, image_size, batch_size):
    image_files = []
    image_labels = []

    for label in labels:
        label_dir = os.path.join(data_dir, label)
        for image_file in os.listdir(label_dir):
            image_files.append(os.path.join(label_dir, image_file))
            image_labels.append(label)

    # Create a mapping from labels to indices
    label_map = {label: idx for idx, label in enumerate(labels)}
    image_labels = [label_map[label] for label in image_labels]

    # Convert to TensorFlow datasets
    dataset = tf.data.Dataset.from_tensor_slices((image_files, image_labels))
    dataset = dataset.map(lambda x, y: load_and_preprocess_image(x, y, image_size))
    dataset = dataset.shuffle(buffer_size=len(image_files))
    dataset = dataset.batch(batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)

    return dataset

def create_and_train_model(base_model, model_name, shape, X_train, X_val, num_classes, labels, log_dir, model_dir,
                           epochs, optimizer_name, learning_rate, step_gamma, alpha, batch_size, patience):
    # Freeze the base model layers
    for layer in base_model.layers:
        layer.trainable = False

    # Add custom layers on top
    x = base_model.output
    x = Flatten()(x)
    x = Dense(1024, activation='relu')(x)
    x = Dropout(0.25)(x)

    x = Dense(512, activation='relu')(x)
    x = Dropout(0.25)(x)

    x = Dense(256, activation='relu')(x)
    x = BatchNormalization()(x)
    x = Dropout(0.25)(x)

    predictions = Dense(num_classes, activation='softmax')(x)  # Use the number of classes
    model = Model(inputs=base_model.input, outputs=predictions)

    # Learning rate schedule
    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
        initial_learning_rate=learning_rate,
        decay_steps=1000,  # Adjust this according to your needs
        decay_rate=step_gamma
    )

    # Select the optimizer
    if optimizer_name.lower() == 'adam':
        optimizer = Adam(learning_rate=lr_schedule)
    elif optimizer_name.lower() == 'sgd':
        optimizer = SGD(learning_rate=lr_schedule, momentum=alpha)  # Example settings for SGD
    else:
        raise ValueError(f"Unsupported optimizer: {optimizer_name}")

    # Compile the model
    model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

    # Set up callbacks
    checkpoint = ModelCheckpoint(os.path.join(model_dir, f'{model_name}_best_model.keras'), 
                                 monitor='val_accuracy', save_best_only=True, save_weights_only=False, 
                                 mode='max', verbose=1)
    early_stopping = EarlyStopping(monitor='val_accuracy', patience=patience, verbose=1)

    # Train the model
    history = model.fit(X_train, epochs=epochs, validation_data=X_val, batch_size=batch_size,
                        callbacks=[checkpoint, early_stopping])

    # Save training logs
    with open(os.path.join(log_dir, f'{model_name}_training.log'), 'w') as f:
        num_epochs = len(history.history['loss'])  # Get the actual number of epochs completed
        for epoch in range(num_epochs):
            f.write(f"Epoch {epoch + 1}, "
                    f"Train Loss: {history.history['loss'][epoch]:.4f}, "
                    f"Train Accuracy: {history.history['accuracy'][epoch]:.4f}, "
                    f"Val Loss: {history.history['val_loss'][epoch]:.4f}, "
                    f"Val Accuracy: {history.history['val_accuracy'][epoch]:.4f}\n")

    # Save labels in the model directory
    with open(os.path.join(model_dir, 'labels.txt'), 'w') as f:
        f.write('\n'.join(labels))

    # Evaluate the model
    test_loss, test_accuracy = model.evaluate(X_val)
    print(f'Test Accuracy for {model_name}: {test_accuracy:.4f}')
    print(f'Test Loss for {model_name}: {test_loss:.4f}')

    # Optionally, save the trained model
    model.save(os.path.join(model_dir, f'{model_name}_final_model.keras'))

def main(base_model_names, shape, data_path, log_dir, model_dir, epochs, optimizer, learning_rate, step_gamma, alpha, batch_size, patience):
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)

    # Extract labels from folder names
    labels = sorted([d for d in os.listdir(os.path.join(data_path, 'train')) if os.path.isdir(os.path.join(data_path, 'train', d))])
    num_classes = len(labels)

    # Load data
    X_train = create_dataset(os.path.join(data_path, 'train'), labels, shape, batch_size)
    X_val = create_dataset(os.path.join(data_path, 'val'), labels, shape, batch_size)

    if not base_model_names:
        print("No base models specified. Exiting.")
        return

    # Define base models
    base_models_dict = {
        model_name: globals()[model_name](weights='imagenet', include_top=False, input_shape=shape)
        for model_name in base_model_names
    }

    for model_name in base_model_names:
        print(f'Training {model_name}...')
        base_model = base_models_dict.get(model_name)
        if base_model is None:
            print(f"Model {model_name} not supported.")
            continue
        create_and_train_model(base_model, model_name, shape, X_train, X_val, num_classes, labels, log_dir, model_dir,
                               epochs, optimizer, learning_rate, step_gamma, alpha, batch_size, patience)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train models using transfer learning")
    parser.add_argument('--base_models', type=str, nargs='+', default=[],
                        help='List of base models to use for training. Leave empty to skip model training.')
    parser.add_argument('--shape', type=int, nargs=3, default=(224, 224, 3), help='Input shape of the images')
    parser.add_argument('--data_path', type=str, required=True, help='Path to the image data')
    parser.add_argument('--log_dir', type=str, required=True, help='Directory to save logs')
    parser.add_argument('--model_dir', type=str, required=True, help='Directory to save models')
    parser.add_argument('--epochs', type=int, default=100, help='Number of epochs to train')
    parser.add_argument('--optimizer', type=str, default='adam', help='Optimizer to use (adam or sgd)')
    parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate')
    parser.add_argument('--step_gamma', type=float, default=0.96, help='Gamma value for step learning rate schedule')
    parser.add_argument('--alpha', type=float, default=0.9, help='Alpha for the optimizer (used for SGD)')
    parser.add_argument('--batch_size', type=int, default=32, help='Batch size for training')
    parser.add_argument('--patience', type=int, default=10, help='Patience for early stopping')

    args = parser.parse_args()
    main(args.base_models, tuple(args.shape), args.data_path, args.log_dir, args.model_dir,
         args.epochs, args.optimizer, args.learning_rate, args.step_gamma, args.alpha, args.batch_size, args.patience)