RAF_DB数据集分类_3

混淆矩阵

这里ECANet太长了，我这里直接利用resnet代替一下，你可以直接替换，然后把权重对应好即可，这只是一个简单的混淆矩阵生成，没有太多美化。

from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import  DataLoader
import torchvision
from torch.nn import init
import math
from torchvision import transforms
import torch.nn.functional as F
val_path = "./RAF-DB/test"
val_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])
])
batch_size = 256
val_data = torchvision.datasets.ImageFolder(val_path, transform=val_transform)
data1_val = DataLoader(val_data, batch_size=batch_size, shuffle=True,drop_last=True)

def plot_confusion_matrix(cm, savename, title='Confusion Matrix'):

    plt.figure(figsize=(12, 8), dpi=100)
    np.set_printoptions(precision=2)

    # 在混淆矩阵中每格的概率值
    ind_array = np.arange(len(classes))
    x, y = np.meshgrid(ind_array, ind_array)
    for x_val, y_val in zip(x.flatten(), y.flatten()):
        c = cm[y_val][x_val]
        if c > 0.001:
            plt.text(x_val, y_val, "%0.2f" % (c,), color='red', fontsize=15, va='center', ha='center')
    
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.binary)
    plt.title(title)
    plt.colorbar()
    xlocations = np.array(range(len(classes)))
    plt.xticks(xlocations, classes, rotation=90)
    plt.yticks(xlocations, classes)
    plt.ylabel('Actual label')
    plt.xlabel('Predict label')
    
    # offset the tick
    tick_marks = np.array(range(len(classes))) + 0.5
    plt.gca().set_xticks(tick_marks, minor=True)
    plt.gca().set_yticks(tick_marks, minor=True)
    plt.gca().xaxis.set_ticks_position('none')
    plt.gca().yaxis.set_ticks_position('none')
    plt.grid(True, which='minor', linestyle='-')
    plt.gcf().subplots_adjust(bottom=0.15)
    
    # show confusion matrix
    # plt.savefig(savename, format='png')
    plt.show()

# classes表示不同类别的名称，比如这有6个类别
classes = ['Anger', 'Disgust', 'Fear', 'Happiness','Neutral', 'Sadness','Surprise']

from torchvision import models
resnet = models.resnet18()
class SKNet(nn.Module):
    def __init__(self, num_class=7):
        super(SKNet, self).__init__()
        self.features = nn.Sequential(*list(resnet.children())[:-2])
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(512, num_class)

    def forward(self, x):
        
        x = self.features(x)
        out = self.avgpool(x)
        out = torch.flatten(out,1)
        out = self.fc(out)
        return out

model_path = "./OneModel/迁移学习/ResNet18.pkl"
sknet = SKNet()
checkpoint = torch.load(model_path,map_location='cpu')
sknet.load_state_dict(checkpoint['model'])
del checkpoint



def evalute_(model,val_loader):
    model.eval()

    for batchidx, (x, label) in enumerate(val_loader):
        with torch.no_grad():
            print(batchidx)
            y1 = model(x)
            _, preds1 = torch.max(F.softmax(y1,dim=1), 1)
            if batchidx!=0:
                y = torch.cat((y,preds1),dim=0)
                labels = torch.cat((labels,label),dim=0)
            else:
                y = preds1
                labels = label
    print(y.shape)
    print(labels.shape)
    assert y.shape == labels.shape
    y = y.numpy()
    
    return y,labels

y_pred,y_true = evalute_(model=sknet,val_loader=data1_val)


# 获取混淆矩阵
cm = confusion_matrix(y_true, y_pred)
cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
plot_confusion_matrix(cm_normalized, './confusion_matrix.png', title='confusion matrix')