Simple Image Classification -TensorFlow

Published by Abhay Rastogi on

Image classification is used for predicting image objects.it assists in many areas like detecting person info, object description and even use to predict skin cancer. In this, we will create over own Ferrari and Lamborgigni dataset and predicted the type of image.

Install dependencies in Python project and download dataset link:-(download_car_dataset)

pip install tensorflow

Create your own image dataset by simply downloading image from the internet.then past it in train_image and test_image folder.as in this we will use car dataset to predict which one is Lamborghini or Ferrari.

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

print(tf.__version__) #Tensorflow version

Set train and test dataset directory location

PATH = 'car_dataset'
train_dir = os.path.join(PATH, 'train_image')
test_dir = os.path.join(PATH, 'test_image')

train_farrari_dir = os.path.join(train_dir, 'farrari')  # directory with our training farrari pictures
train_lamborghini_dir = os.path.join(train_dir, 'lamborghini')  # directory with our training lamborghini pictures

test_farrari_dir = os.path.join(test_dir, 'farrari')  # directory with our test farrari pictures
test_lamborghini_dir = os.path.join(test_dir, 'lamborghini')  # directory with our test lamborghini pictures

Compute the total length of train and test images.

num_farrari_tr = len(os.listdir(train_farrari_dir))
num_lamborghini_tr = len(os.listdir(train_lamborghini_dir))

num_farrari_val = len(os.listdir(test_farrari_dir))
num_lamborghini_val = len(os.listdir(test_lamborghini_dir))

total_train = num_farrari_tr + num_lamborghini_tr # total training set
total_val = num_farrari_val + num_lamborghini_val # total test set

Formate your image data to be read into gride format then convert it into tensor and resize all the images into the same height and width.

batch_size = 128
epochs = 15
IMG_HEIGHT = 150
IMG_WIDTH = 150

train_image_generator = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255) # Generator for our training data
test_image_generator =tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255) # Generator for our test data

train_data_gen = train_image_generator.flow_from_directory(batch_size=batch_size,
                                                           directory=train_dir,
                                                           shuffle=True,
                                                           target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                           class_mode='binary')

test_data_gen = test_image_generator.flow_from_directory(batch_size=batch_size,
                                                              directory=test_dir,
                                                              target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                              class_mode='binary')

print(train_data_gen.class_indices) # gives index of each lable 
#output: {'farrari': 0, 'lamborghini': 1}

Create a model of 512 units connected layer with relu activation function and compile it adam optimizer with BinaryCrossentropy loss.

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu', input_shape=(IMG_HEIGHT, IMG_WIDTH ,3)),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation='relu'),
    tf.keras.layers.Dense(1)
])
model.compile(optimizer='adam',
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.summary()

history = model.fit_generator(
    train_data_gen,
    steps_per_epoch= total_train,
    epochs=epochs,
    validation_data= test_data_gen,
    validation_steps= total_val
)
model.save("car_model.h5") # save the model

Predict Image

import tensorflow as tf
import  numpy as np

model=tf.keras.models.load_model('car_model.h5')#path of pre-trained model
categories = ["FARRARI","LAMBORGHINI"]
urlimage=tf.keras.preprocessing.image.load_img('farrari01.jpg',target_size=(150,150)) #path of random car image
urlimage = tf.keras.preprocessing.image.img_to_array(urlimage)
urlimage = np.expand_dims(urlimage,axis=0)
urlimage = urlimage/255
result= model.predict_classes(urlimage)[0][0]
print(result)
print(categories[result])
Categories: ML

0 Comments

Leave a Reply

Avatar placeholder

Your email address will not be published.