evaluate.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import os
import yaml
import models as models
from dataset_preparation import MalwareDetect2Dataset, OpcodeDataset
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import numpy as np
from sklearn.metrics import precision_score, recall_score, f1_score
import pandas as pd


def evaluate_performace(model_name, model, data_loader):
    correct = 0
    total = 0
    all_labels = []
    all_predections = []

    with torch.no_grad():
        for images, labels in data_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            all_labels.extend(labels.cpu().numpy())
            all_predections.extend(predicted.cpu().numpy())
    
    accuracy = 100 * correct / total
    precision = 100 * precision_score(all_labels, all_predections, average='micro')
    recall = 100 * recall_score(all_labels, all_predections, average='micro')
    f1 = 100 * f1_score(all_labels, all_predections, average='micro')

    performance_df = pd.DataFrame({
        'Classifier': [model_name],
        'Accuracy': [accuracy],
        'Precision': [precision],
        'Recall': [recall],
        'F1': [f1],
        'AUC': [None]
    })

    print(performance_df)

def generate_confusion_matrix(model, test_loader, classes):
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)

            # Get predictions.
            outputs = model(images)
            _, preds = torch.max(outputs, 1)

            # Append predictions and true labels.
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    # Generate confusion matrix.
    cm = confusion_matrix(all_labels, all_preds)

    # Display confusion matrix.
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=classes)
    disp.plot(cmap=plt.cm.Blues)
    plt.title("Confusion Matrix LinearClassifierTFIDF")
    plt.yticks(rotation=45)
    plt.savefig("./artefacts/confusion_matrices/cm_linearclassifiertfidf.png")

    return cm

# Configure device.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# checkpoint = torch.load('./models/cnn_3mpdofc_1layer_net.pth')
checkpoint = torch.load('./models/linear_classifier.pth')

'''Extract the hyperparameters and training history.'''
input_size = checkpoint['input_size']
# hidden_size = checkpoint['hidden_size']
# hidden_size1 = checkpoint['hidden_size1']
# hidden_size2 = checkpoint['hidden_size2']
# num_classes = checkpoint['num_classes']
# input_channels = checkpoint['input_channels']
num_epochs = checkpoint['num_epochs']
batch_size = checkpoint['batch_size']
learning_rate = checkpoint['learning_rate']
history = checkpoint['history']

# model = models.CNN3MPDOFC1LayerNet(1, hidden_size, 5).to(device)
model = models.LinearClassifier(393, 5).to(device)

# Load the model's state.
model.load_state_dict(checkpoint['model_state_dict'])

# Create image transformations.
transform_fc = transforms.Compose([
    transforms.ToTensor(),
])
transform_cnn = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

with open("config.yaml") as f:
    config = yaml.safe_load(f)

data = config["data"]["opcode_data_tfidf"]
opcode_frequency_data_zipped = np.load(data)
X_train = opcode_frequency_data_zipped['X_train']
y_train = opcode_frequency_data_zipped['y_train']
X_test = opcode_frequency_data_zipped['X_test']
y_test = opcode_frequency_data_zipped['y_test']

# train_dir = "./data/malware_detect2_simhash_1616_images/train"
# test_dir = "./data/malware_detect2_simhash_1616_images/test"

# Train dataset and DataLoader
# train_dataset = MalwareDetect2Dataset(data_dir=train_dir, transform=transform_cnn)
# train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
train_dataset = OpcodeDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)

# Test dataset and DataLoader
# test_dataset = MalwareDetect2Dataset(data_dir=test_dir, transform=transform_cnn)
# test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
test_dataset = OpcodeDataset(X_test, y_test)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

model.eval()
# evaluate_performace("CNN3MPDOFC1", model, train_loader)
# evaluate_performace("CNN3MPDOFC1", model, test_loader)
evaluate_performace("LinearClassifierTFIDF", model, train_loader)
evaluate_performace("LinearClassifierTFIDF", model, test_loader)

# Given with the assumtion that the classes are ordered alphabetically.
classes = ["Locker", "Mediyes", "Winwebsec", "Zbot", "Zeroaccess"]
cm = generate_confusion_matrix(model, test_loader, classes)