Search code examples
pythontensorflowconfusion-matrix

Find confusion matrix of image classification in Tensorflow


I am training a model to classify images into 2 classes following this tutorial: https://www.tensorflow.org/tutorials/images/classification

After model.fit(), I want to evaluate the accuracy of the model's prediction using a test set which contains images that weren't included in the training or validation sets. The test set contains 2 folders which contain the images of the corresponding class.

├── test_data/
│   ├── class1/
│   ├── class2/

And I want to find the recall, precision and accuracy of each class using confusion matrix. However, I am new to deep learning and Tensorflow. I have no idea how to get the confusion matrix for each class. I am also not sure whether the way I pass images to the model is correct.

The following is my current implementation on using the model to predict on new data.

# get the list of class names in the training set
train_class_names = train_ds.class_names

# load the test data
test_data_dir = pathlib.Path('test_data/')
test_data_list = list(test_data_dir.glob('*/*.jpg'))

test_ds = tf.keras.utils.image_dataset_from_directory(
  test_data_dir,
  image_size=(img_height, img_width),
  batch_size=batch_size)

predicted_img = []

# for every image in the test_data folder, pass it to the model to predict its class
for path in test_data_list:
    img = tf.keras.utils.load_img(
        path, target_size=(img_height, img_width)
    )
    img_array = tf.keras.utils.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)
    
    test_class_name = path.parent.name

    predictions = model.predict(img_array)
    score = tf.nn.softmax(predictions[0])
    
    # append the image, predicted class and actual class to a list 
    # so that I can print them out to see if the prediction is correct
    predicted_img.append([img, train_class_names[np.argmax(score)], test_class_name])

Solution

  • try this

    from sklearn.metrics import confusion_matrix, classification_report
    import seaborn as sns
    sns.set_style('darkgrid')
    classes=test_ds.class_names # ordered list of class names
    ytrue=[]
    for images, label in test_ds:   
        for e in label:
            ytrue.append(classes[e]) # list of class names associated with each image file in test dataset 
    ypred=[]
    errors=0
    count=0
    preds=model.predict(test_ds, verbose=1) # predict on the test data
    for i, p in enumerate(preds):
        count +=1
        index=np.argmax(p) # get index of prediction with highest probability
        klass=classes
    [index] 
        ypred.append(klass)  
        if klass != ytrue[i]:
            errors +=1
    acc= (count-errors)* 100/count
    msg=f'there were {count-errors} correct predictions in {count} tests for an accuracy of {acc:6.2f} % '
    print(msg) 
    ypred=np.array(ypred)
    ytrue=np.array(ytrue)
    if len(classes)<= 30: # if more than 30 classes plot is not useful to cramed
            # create a confusion matrix 
            cm = confusion_matrix(y_true, y_pred )        
            length=len(classes)
            if length<8:
                fig_width=8
                fig_height=8
            else:
                fig_width= int(length * .5)
                fig_height= int(length * .5)
            plt.figure(figsize=(fig_width, fig_height))
            sns.heatmap(cm, annot=True, vmin=0, fmt='g', cmap='Blues', cbar=False)       
            plt.xticks(np.arange(length)+.5, classes, rotation= 90)
            plt.yticks(np.arange(length)+.5, classes, rotation=0)
            plt.xlabel("Predicted")
            plt.ylabel("Actual")
            plt.title("Confusion Matrix")
            plt.show()
    clr = classification_report(ytrue, ypred, target_names=class_names)
    print("Classification Report:\n----------------------\n", clr)