TensorFlow Binary Image Classification: Predict Probability of each class for each image in data set

You don't have to make your model output "Good" and "Bad" as labels, instead, you could output probabilities for each one independently, in other words, the probability of the image being good and the probability of the image being bad. Make the size of the output of your last layer to be 2. So your model will now output a 2-dimensional vector such that [1.0, 0.0] means 100% good and 0% bad and [0.0, 1.0] means 0% good and 100% bad. Use binary cross-entropy as your loss function for training. Of course, you have to label your training data similarly, so if you have a good training example, label it as [1.0, 0.0] since you are 100% sure that it is good and if you have a bad training example label it as [0.0, 1.0] since you are also 100% sure that it is a bad example.

The reason I told you to use binary-cross entropy as a loss function is so that the model will learn to output opposing probabilities for the components of the 2-d vector output. So if it is a good image, the first component will be high and the second component will be low and vice-versa if it is a bad image. Also, after training, when making predictions, you only take the highest probability of the two, if the higher probability is the first one, then it is a "Good" image and you use that probability only.

In case someone is looking for an answer, below is the python code for model generation

Some of the points to be noted here are

1. Input Image shape is 360x360x3
2. Activation Function for last layer is "softmax" instead of "sigmoid"
3. Loss function is "sparse_categorical_crossentropy" instead of "binary_crossentropy"
4. Shape of output is 2 instead of 1

Please note #2, #3 and #4, even though I am trying to come up with a model for Binary Image Classification. My ultimate aim was to convert this model to TensorFlow Lite version and use TensorFlow Lite model in Android Application.

Earlier, when I was using "sigmoid" for last layer and "binary_crossentropy" as loss function, the output shape of last layer cannot be greater than 1.

As a result when I was using the Lite model generated out of that TensorFlow Model in Android application, I was getting an error which is mentioned below

"Cannot find an axis to label. A valid axis to label should have size larger than 1"

With the changes mentioned in #2, #3, and #4, the Lite Model generated works fine in Android.

import tensorflow as tf
import matplotlib.pyplot as plt
import cv2
import os
import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.preprocessing import image
from tensorflow.keras.optimizers import RMSprop


print("version")
print(tf.__version__)

train = ImageDataGenerator(rescale=1/255)
validation = ImageDataGenerator(rescale=1/255)

input_shape = 360
train_dataset = train.flow_from_directory('container_images/train/',
                                          target_size=(input_shape,input_shape),
                                          batch_size=3,
                                          classes=['good', 'bad'],
                                          class_mode='binary')

validation_dataset = train.flow_from_directory('container_images/validation/',
                                          target_size=(input_shape,input_shape),
                                          batch_size=3,
                                          classes=['good', 'bad'],
                                          class_mode='binary')

print(train_dataset.class_indices)
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(16, (3,3), activation='relu', input_shape=(input_shape, input_shape, 3)),
    tf.keras.layers.MaxPool2D(2,2),
    #
    tf.keras.layers.Conv2D(32, (3,3), activation='relu'),
    tf.keras.layers.MaxPool2D(2,2),
    #
    tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
    tf.keras.layers.MaxPool2D(2,2),
    ##
    tf.keras.layers.Flatten(),
    ##
    tf.keras.layers.Dense(512, activation='relu'),
    ##
    tf.keras.layers.Dense(2, activation='softmax')
])

model.compile(loss='sparse_categorical_crossentropy',
              optimizer=RMSprop(lr=0.001),
              metrics=['accuracy'])
model_fit = model.fit(train_dataset,
                      steps_per_epoch=3,
                      epochs=30,
                      validation_data=validation_dataset)