Animal Classifier using Transfer Learning with MobileNet V2

45 Views Asked by At

My code below compiles perfectly fine. The data set I have gave to the program is simple, its just couple of animals within some folders and it understands what animal is it.

The code is:

It compiles find. I just want to run the code and give to it a specific .jpg image as an input. How can I do it?

1

There are 1 best solutions below

1
YokiT On
    import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import numpy as np
import os
import random
from PIL import Image

# Set the paths to your dataset and animal names file
data_dir = '/Users/admin/PycharmProjects/ML/animals2'
animal_names_file = '/Users/admin/PycharmProjects/ML/animalNames2.txt'

# Load animal names from the text file
with open(animal_names_file, 'r') as file:
    animal_names = file.read().splitlines()

# Set the parameters for the model and training
image_size = (224, 224)
batch_size = 32
epochs = 10
num_classes = len(animal_names)

# Preprocess and augment the data
data_generator = ImageDataGenerator(
    rescale=1./255,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    validation_split=0.2
)

# Generate training and validation datasets
train_generator = data_generator.flow_from_directory(
    data_dir,
    target_size=image_size,
    batch_size=batch_size,
    class_mode='categorical',
    subset='training'
)

validation_generator = data_generator.flow_from_directory(
    data_dir,
    target_size=image_size,
    batch_size=batch_size,
    class_mode='categorical',
    subset='validation'
)

# Build the model
base_model = tf.keras.applications.MobileNetV2(input_shape=(224, 224, 3), include_top=False, weights='/Users/admin/PycharmProjects/ML/mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_224_no_top.h5')
base_model.trainable = False

model = tf.keras.models.Sequential([
    base_model,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(num_classes, activation='softmax')
])

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

# Train the model
model.fit(
    train_generator,
    steps_per_epoch=train_generator.n // batch_size,
    epochs=epochs,
    validation_data=validation_generator,
    validation_steps=validation_generator.n // batch_size
)

# Save the trained model
model.save('animal_classifier_model.keras')
# Select a random image from each sub-folder
animal_folder = os.path.join(data_dir)
animal_subfolders = [subfolder for subfolder in os.listdir(animal_folder) if os.path.isdir(os.path.join(animal_folder, subfolder))]

# Load the saved model outside the loop
loaded_model = tf.keras.models.load_model('animal_classifier_model.h5')

# Define a function for making predictions
def predict_image(image_path):
    # Load and preprocess the test image
    test_image = Image.open(image_path)
    test_image = test_image.resize(image_size)
    test_image = np.array(test_image) / 255.0
    test_image = np.expand_dims(test_image, axis=0)

    # Make predictions using the loaded model
    predictions = loaded_model.predict(test_image)
    predicted_class = np.argmax(predictions)

    return predicted_class

# Make predictions for each sub-folder
for subfolder in animal_subfolders:
    subfolder_path = os.path.join(animal_folder, subfolder)
    animal_images = os.listdir(subfolder_path)
    random_image = random.choice(animal_images)
    image_path = os.path.join(subfolder_path, random_image)

    # Call the predict_image function to makepredictions for the current image path
    predicted_class = predict_image(image_path)

    # Get the predicted animal name from the animal_names list
    predicted_animal = animal_names[predicted_class]

    print("Predicted animal for", subfolder, ":", predicted_animal)