Luis Jose Mendez
Person 🧑🏻‍🦱 - Dog 🐶 - Cat 🐈 Multi Classification Model

Person 🧑🏻‍🦱 - Dog 🐶 - Cat 🐈 Multi Classification Model

Luis Jose Mendez 's photo
Luis Jose Mendez
·

2 min read

Table of contents

Multi-Classification Person-Dog-Cat Model with Transfer Learning.

The Model predicts if a given image is a Dog, Cat or Person.

It was trained using Transfer Learning, with the VGG-16 Model. The last output layer was removed and then updated with an output layer of three neurons, with a softmax activation function.

The Model was retrained with 1000 images from each class.

Modules

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
from tensorflow import keras
import numpy as np

Loading the Pre-Trained VGG-16 Model from TensorFlow Library

vgg_model = tf.keras.applications.vgg16.VGG16()
print(type(vgg_model))

preprocess_input = tf.keras.applications.vgg16.preprocess_input

vgg_model.summary()

Loading Train, Validation and Test Images

train_data = keras.preprocessing.image.ImageDataGenerator(preprocessing_function=preprocess_input)

validation_data = keras.preprocessing.image.ImageDataGenerator(preprocessing_function=preprocess_input)

test_data = keras.preprocessing.image.ImageDataGenerator(preprocessing_function=preprocess_input)

BATCH_SIZE = 32
PIC_SIZE = 224

train = train_data.flow_from_directory("./train", target_size=(PIC_SIZE,PIC_SIZE), batch_size=BATCH_SIZE, class_mode="categorical", color_mode='rgb', shuffle=True)
validation = test_data.flow_from_directory("./validation", target_size=(PIC_SIZE,PIC_SIZE), batch_size=BATCH_SIZE, class_mode="categorical", color_mode='rgb', shuffle=True)
test = test_data.flow_from_directory("./test", target_size=(PIC_SIZE,PIC_SIZE), batch_size=BATCH_SIZE, class_mode="categorical", color_mode='rgb', shuffle=True)

Adding Every Single Layer of the VGG-16 Model Except the Last One

model = keras.models.Sequential()
for layer in vgg_model.layers[0:-1]:
    model.add(layer)

Deactivating Every Trainable Layer, so The New Last One Layer is Going to be Retrained

for layer in model.layers:
    layer.trainable = False

Adding the last Layer

model.add(tf.keras.layers.Dense(3, activation='softmax'))

Compiling and Retraining the Model

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
    loss=tf.keras.losses.CategoricalCrossentropy(),
    metrics=['acc']
)

history = model.fit(train,   
    epochs=10, validation_data=validation)

The Model has a very high percentage Accuracy, between 95%-98%. At the bottom of the page are going to be all the datasets links and the GitHub repository with all the code.

Testing the Model

test_img = tf.keras.preprocessing.image.load_img("path", target_size=(224, 224), color_mode='rgb')

plt.imshow(test_img)

X = tf.keras.preprocessing.image.img_to_array(test_img)
X = np.expand_dims(X, axis = 0)

prediction = np.vstack([X])

result = model.predict(prediction)

arg_max_result = np.argmax(result)

print(result)

if arg_max_result == 0 :
    print("Cat")
elif arg_max_result == 1 :
    print("Dog")
elif arg_max_result == 2 :
    print("Human")

Saving the Model

model.save("dog_cat_person.h5")

Check-it out

Test the Model by yourself by running the main.py file, built with Streamlit. Remember first to run the model jupyter notebook file and save it.

streamlit run main.py

Resources:

Transfer learningTensorFlowmulti-classification