model_definition.py

from vit_keras import vit
from tensorflow.keras.layers import Dense, Dropout, Flatten
from tensorflow.keras.models import Model
import tensorflow as tf

# Check if GPU or CPU will be used
device_name = "GPU" if tf.test.is_gpu_available() else "CPU"
print(f"Using {device_name} for processing")

def create_model():
    # Load a pre-trained Vision Transformer model
    base_model = vit.vit_b16(
        image_size=224,
        pretrained=True,
        include_top=False,
        pretrained_top=False
    )

    # Add classification head
    x = base_model.output
    x = Flatten()(x)  # Flatten the feature map
    x = Dense(128, activation='relu')(x)
    x = Dropout(0.5)(x)  # Add dropout for regularization
    output = Dense(1, activation='sigmoid')(x)  # Binary classification

    model = Model(inputs=base_model.input, outputs=output)

    # Freeze the base model layers for transfer learning
    base_model.trainable = False

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

    return model