Pythonで混同行列と分類レポートの生成

機械学習では、モデルの評価のために混同行列（Confusion Matrix）や分類レポート（Classification Report）を可視化することがよくあります。

そこで、以下では、Pythonの scikit-learn、matplotlib、seaborn を使用して、これらを画像として保存します。

混同行列
#

まず、混同行列を可視化する関数を定義します。

python

import matplotlib.figure
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from sklearn.metrics import confusion_matrix


def plot_confusion_matrix(y_true: list, y_pred: list, label_list: list, title: str = "Confusion matrix") -> matplotlib.figure.Figure:
    conf_mat = confusion_matrix(y_true, y_pred, labels=label_list)
    print(conf_mat)
    fig, ax = plt.subplots()
    sns.heatmap(conf_mat, annot=True, fmt="d", xticklabels=label_list, yticklabels=label_list)
    plt.ylabel("True")
    plt.xlabel("Pred")
    plt.title(title)

    plt.tight_layout()

    return fig

Classification report
#

次に、分類レポートをヒートマップとして可視化する関数を定義します。

python

import matplotlib.figure
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from sklearn.metrics import classification_report


def plot_classification_report(y_true: list, y_pred: list, label_list: list, title: str = "Classification report") -> matplotlib.figure.Figure:
    report = classification_report(y_true, y_pred, labels=label_list, output_dict=True)
    df_report = pd.DataFrame(report).transpose()

    fig, ax = plt.subplots(figsize=(10, 7))
    sns.heatmap(df_report.iloc[:-1, :-1], annot=True)
    plt.title(title)

    plt.tight_layout()

    return fig

使用方法
#

定義した関数を使って、以下のように混同行列と分類レポートの画像を作成・保存します。

python

y_true = ["cat", "dog", "bird", "cat", "bird", "bird"]
y_pred = ["cat", "cat", "bird", "cat", "cat", "bird"]
label_list = ["dog", "cat", "bird"]

fig1 = plot_confusion_matrix(y_true, y_pred, label_list, title="My Confusion Matrix")
fig2 = plot_classification_report(y_true, y_pred, label_list, title="My Classification Report")
fig1.savefig('confusion_matrix.png')
fig2.savefig('classification_report.png')