[FIXED] how to build cnn model for defect detection with three classes and test one image on it

Issue

I built a CNN model to detect two kinds of defects on an image. These classes are 'big' and 'small' and the accuracy is really good. The architecture of my model is in the below:

inputs = tf.keras.Input(shape=(120, 120, 3))
x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu')(inputs)
x = tf.keras.layers.MaxPool2D(pool_size=(2, 2))(x)
x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu')(x)
x = tf.keras.layers.MaxPool2D(pool_size=(2, 2))(x)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)

model = tf.keras.Model(inputs=inputs, outputs=outputs)

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

print(model.summary())

history = model.fit(
    train_data,
    validation_data=val_data,
    epochs=100,
    callbacks=[
        tf.keras.callbacks.EarlyStopping(
            monitor='val_loss',
            patience=3,
            restore_best_weights=True
        )
    ]
)

Now, I want to use this CNN model for multi classes and the classes would be 'big', 'small', 'other'. I have the data set, but I don't know how to change the model for three classes. Also, at the end I want to test one image to my CNN model and get the label if the inserted image is big, small or other, but I don't know how.

Solution

Try this:

model = tf.keras.Sequential()
model.add(tf.keras.layers.Conv2D(input_shape = (120, 120, 3), filters=16, kernel_size=(3, 3), activation='relu'))
model.add(tf.keras.layers.MaxPool2D(pool_size=(2, 2)))
model.add(tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation='relu'))
model.add(tf.keras.layers.MaxPool2D(pool_size=(2, 2)))
model.add(tf.keras.layers.GlobalAveragePooling2D())
model.add(tf.keras.layers.Dense(3, activation='softmax'))

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

to predict, you can use this code:

from PIL import Image
import numpy as np
from skimage import transform

def load(filename):
    np_image = Image.open(filename)
    np_image = np.array(np_image).astype('float32')/255
    np_image = transform.resize(np_image, (120, 120, 3))
    np_image = np.expand_dims(np_image, axis=0)
    return np_image


folder_path = 'Dataset/test/4.jpg'
image = load(folder_path)
pred = model.predict_classes(image)
pred.tolist()[0]

Answered By - Adarsh Wase