How to Create a Confusion Matrix in Python

A confusion matrix is a specific table layout that allows visualization of the performance of an algorithm. In classification tasks, it compares the predicted labels against the actual labels (ground truth), breaking down results into true positives, true negatives, false positives, and false negatives.

This guide explains how to implement a confusion matrix function from scratch using Python and NumPy, processing raw probability predictions into a readable grid.

Understanding the Inputs

To build a matrix, we need three pieces of data:

1. Labels: The names of the classes (e.g., ["Cat", "Dog", "Bird"]).
2. Predictions (preds): A list of probability vectors. For example, [0.1, 0.8, 0.1] implies the model is 80% sure it's the second class ("Dog").
3. Ground Truth: The actual correct labels for the data.

Our goal is to map the highest probability in the prediction to a row, and the actual label to a column, then increment the counter at that intersection.

Step 1: Setting Up the Environment

We rely on the numpy library to efficiently find the index of the maximum value in a prediction list (the class the model "chose").

pip install numpy

Start your file confusion_matrix.py with the necessary imports:

import numpy as np
from typing import List

Step 2: Implementing the Function

The core logic requires iterating through every prediction, converting the probability distribution into a concrete class index, and matching it against the truth.

The Algorithm

def confusion_matrix(
    labels: List[str], 
    preds: List[List[float]], 
    ground_truth: List[str]
) -> List[List[int]]:
    """
    Computes the confusion matrix.
    Rows = Predicted Class
    Columns = True Class
    """
    
    num_classes = len(labels)
    matrix = [[0 for _ in range(num_classes)] for _ in range(num_classes)]

    for pred_probs, truth_label in zip(preds, ground_truth):
        
        pred_index = np.argmax(pred_probs)
        
        try:
            truth_index = labels.index(truth_label)
        except ValueError:
            continue # Skip if label is not found in our definitions

        matrix[pred_index][truth_index] += 1

    return matrix

Step 3: Testing with Example Data

To ensure our logic holds, we test with a small dataset where we can manually verify the results.

Test Scenario

• Labels: Python (0), Java (1), C++ (2)
• Data: 5 samples.
• Logic: If the model predicts "Python" but the truth is "Java", we add 1 to Row 0, Column 1.

if __name__ == "__main__":
    labels = ["Python", "Java", "C++"]
    
    preds = [
        [0.66, 0.22, 0.11], # Max is index 0 (Python)
        [0.34, 0.05, 0.60], # Max is index 2 (C++)
        [0.47, 0.26, 0.26], # Max is index 0 (Python)
        [0.76, 0.15, 0.08], # Max is index 0 (Python)
        [0.05, 0.95, 0.00], # Max is index 1 (Java)
    ]
    
    ground_truth = ["Python", "C++", "Java", "C++", "Java"]

    result = confusion_matrix(labels, preds, ground_truth)
    
    print("Confusion Matrix:")
    for row in result:
        print(row)

Output:

Confusion Matrix:
[1, 1, 1]
[0, 1, 0]
[0, 0, 1]

Interpreting the Output

• Row 0 (Predicted Python): [1, 1, 1]
  • 1 instance was actually Python (Correct).
  • 1 instance was actually Java (Incorrect).
  • 1 instance was actually C++ (Incorrect).
  • Insight: The model is over-predicting "Python".
• Row 1 (Predicted Java): [0, 1, 0]
  • 1 instance was actually Java (Correct).
• Row 2 (Predicted C++): [0, 0, 1]
  • 1 instance was actually C++ (Correct).

Conclusion

By implementing the confusion matrix manually, you gain insight into how classification accuracy is derived.

1. np.argmax converts probabilities to a hard class prediction.
2. zip allows you to process predictions and truth in parallel.
3. Matrix indexing (matrix[pred][truth]) accumulates the results for analysis.