图像分类模型评估之用python绘制混淆矩阵confusion_matrix

目录

混淆矩阵：是用于评估分类模型性能的一种常用工具，可以用于计算分类准确率、精确率、召回率和F1-score等指标。生成混淆矩阵需要将模型对测试集的预测结果与真实标签进行比对，然后统计每个类别被正确预测的数量以及被错误预测的数量，最终将这些数据组织成矩阵的形式。

python实现混淆矩阵代码：

训练完成后，生成混淆矩阵！！！！

ImageNet数据格式，生成混淆矩阵！！！！

其中，data_path是数据集路径，model_path是模型路径，需要根据实际情况进行修改。代码中使用了ImageFolder直接导入数据集，不需要重新定义类。导入时，将数据集根目录和数据增强方法传入ImageFolder中即可。最后，生成混淆矩阵并保存为CSV文件。

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from sklearn.metrics import confusion_matrix
import pandas as pd

# 设置设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 定义数据增强
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
data_path = "path/to/dataset"
dataset = datasets.ImageFolder(root=data_path, transform=transform)

# 加载模型
model_path = "path/to/model.pth"
model = torch.load(model_path)
model.to(device)
model.eval()

# 获取预测结果和标签
labels = []
preds = []
for inputs, targets in dataset:
    inputs = inputs.unsqueeze(0).to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    _, predicted = torch.max(outputs.data, 1)
    labels.append(targets.item())
    preds.append(predicted.item())

# 生成混淆矩阵
cm = confusion_matrix(labels, preds)
classes = dataset.classes
cm_df = pd.DataFrame(cm, index=classes, columns=classes)

# 保存为CSV文件
cm_df.to_csv("confusion_matrix.csv")
print("Confusion matrix saved as confusion_matrix.csv")

非 ImageNet数据格式，定义导入数据的类名和方法！！！！

导入数据的类和函数定义代码：

import os
import numpy as np
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image

class CustomDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.img_files = os.listdir(data_dir)

    def __len__(self):
        return len(self.img_files)

    def __getitem__(self, index):
        img_path = os.path.join(self.data_dir, self.img_files[index])
        img = Image.open(img_path).convert('RGB')
        label = self.get_label_from_filename(self.img_files[index])
        
        if self.transform:
            img = self.transform(img)

        return img, label

    def get_label_from_filename(self, filename):
        label = filename.split('.')[0] # 假设文件名为"label.image_id.jpg"格式
        label = label.split('_')[0] # 仅保留label信息
        return int(label)

# 加载数据集并进行预处理
data_dir = "your_data_dir"
transform = transforms.Compose([
    transforms.Resize((224, 224)), # 图像大小调整为224x224
    transforms.ToTensor(), # 将图像转换为Tensor格式，并将像素值缩放到[0, 1]
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 图像标准化
])
dataset = CustomDataset(data_dir, transform=transform)

非imageNet数据格式，完成生成混淆矩阵程序代码！！！！

注意点：

1.自定义数据通过 “_” 来进行获取图片的类别 label 标签值，所以你的命名中是否包含了标签值，如果没有标签值，还是需要自行修改的，当然有的话，也需要思考一下，标签值放置的位置在spilt后列表的那个位置！

2.对了，那个后面部分需要缩进，哈哈，python代码简洁，但缩进问题就很抽象，从

model = torch.load(‘model.pth’)  开始直接缩进！！！

import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix
import pandas as pd
import numpy as np


# 自定义数据集类
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.img_list = os.listdir(root_dir)

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.img_list[idx])
        image = Image.open(img_name).convert('RGB')
        if self.transform:
            image = self.transform(image)
        label = int(self.img_list[idx].split('_')[0])  # 根据文件名获取标签
        return image, label

 # 加载模型
model = torch.load('model.pth')

# 定义数据集
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
dataset = MyDataset(root='path/to/dataset', transform=transform)
dataloader = DataLoader(dataset, batch_size=16, shuffle=False)

# 预测结果和真实标签
y_pred = []
y_true = []
with torch.no_grad():
    for images, labels in dataloader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        y_pred.extend(predicted.cpu().numpy())
        y_true.extend(labels.cpu().numpy())

# 生成混淆矩阵
cm = confusion_matrix(y_true, y_pred)

# 将混淆矩阵保存为CSV文件
pd.DataFrame(cm).to_csv('confusion_matrix.csv', index=False, header=False)

# 打印混淆矩阵
print(cm)