src.py

# Import necessary libraries
import tensorflow as tf
from tensorflow.keras import layers, models, datasets
import matplotlib.pyplot as plt

# Load and preprocess the dataset (using CIFAR-10 as an example)
(train_images, train_labels), (test_images, test_labels) = datasets.cifar10.load_data()  # Load CIFAR-10 dataset
train_images, test_images = train_images / 255.0, test_images / 255.0  # Normalize pixel values to be between 0 and 1

# Define the model (using MobileNetV2 as an example)
base_model = tf.keras.applications.MobileNetV2(input_shape=(32, 32, 3),  # Use MobileNetV2 as the base model
                                               include_top=False,  # Do not include the top layer
                                               weights='imagenet')  # Use pre-trained weights from ImageNet
base_model.trainable = False  # Freeze the base model

# Create a new model with the base model and additional layers
model = models.Sequential([
    base_model,  # Add the base model
    layers.GlobalAveragePooling2D(),  # Add a global average pooling layer
    layers.Dense(10, activation='softmax')  # Add a dense layer with 10 units (for 10 classes) and softmax activation
])

# Compile the model
model.compile(optimizer='adam',  # Use Adam optimizer
              loss='sparse_categorical_crossentropy',  # Use sparse categorical crossentropy loss
              metrics=['accuracy'])  # Track accuracy as a metric

# Train the model
model.fit(train_images, train_labels, epochs=5, validation_data=(test_images, test_labels))  # Train the model for 5 epochs

# Evaluate the model
test_loss, test_acc = model.evaluate(test_images, test_labels)  # Evaluate the model on the test data
print(f"Test accuracy: {test_acc}")  # Print the test accuracy

# Make predictions
predictions = model.predict(test_images)  # Make predictions on the test images

# Visualize results (optional)
plt.figure(figsize=(10, 10))  # Create a figure for plotting
for i in range(25):  # Loop through the first 25 test images
    plt.subplot(5, 5, i+1)  # Create a subplot for each image
    plt.xticks([])  # Remove x-axis ticks
    plt.yticks([])  # Remove y-axis ticks
    plt.imshow(test_images[i])  # Display the image
    predicted_label = tf.argmax(predictions[i])  # Get the predicted label
    true_label = test_labels[i][0]  # Get the true label
    color = 'green' if predicted_label == true_label else 'red'  # Set the color to green if correct, red if incorrect
    plt.xlabel(f"{predicted_label} ({true_label})", color=color)  # Display the predicted and true labels

plt.show()  # Show the plot