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

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):


### 最后

不知道你们用的什么环境，我一般都是用的Python3.6环境和pycharm解释器，没有软件，或者没有资料，没人解答问题，都可以免费领取（包括今天的代码），过几天我还会做个视频教程出来，有需要也可以领取~  

给大家准备的学习资料包括但不限于：  

Python 环境、pycharm编辑器/永久激活/翻译插件  

python 零基础视频教程  

Python 界面开发实战教程  

Python 爬虫实战教程  

Python 数据分析实战教程  

python 游戏开发实战教程  

Python 电子书100本  

Python 学习路线规划

![](https://img-blog.csdnimg.cn/d29631674929476f9c3b30f7ff58dff0.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2ZlaTM0Nzc5NTc5MA==,size_16,color_FFFFFF,t_70)




**网上学习资料一大堆，但如果学到的知识不成体系，遇到问题时只是浅尝辄止，不再深入研究，那么很难做到真正的技术提升。**

**[需要这份系统化学习资料的朋友，可以戳这里无偿获取](https://bbs.csdn.net/topics/618317507)**

**一个人可以走的很快，但一群人才能走的更远！不论你是正从事IT行业的老鸟或是对IT行业感兴趣的新人，都欢迎加入我们的的圈子（技术交流、学习资源、职场吐槽、大厂内推、面试辅导），让我们一起学习成长！**