MalwareDetectionSystem/evaluate.py at main · ThaRealJozef/MalwareDetectionSystem · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import os
import sys
import numpy as np
from tensorflow.keras.models import load_model
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt

def evaluate_test_set(
    model_path="src/models/best_model.h5",
    data_path="data/dataset.npz"
):
    # 1. Load model and data
    model = load_model(model_path)
    data = np.load(data_path)
    X_test, y_test = data["X_test"], data["y_test"]

    # 2. Predict
    preds = model.predict(X_test).ravel()
    y_pred = (preds >= 0.5).astype(int)

    # 3. Determine which classes are present
    classes = np.unique(np.concatenate([y_test, y_pred]))
    # Map 0->benign, 1->malware
    label_names = ["benign", "malware"]
    names = [label_names[int(c)] for c in classes]

    if len(classes) < 2:
        print(f"⚠️ Only one class ({names[0]}) present. Reports will be for that class only.\n")

    # 4. Classification report
    print("CLASSIFICATION REPORT:\n")
    print(classification_report(
        y_test,
        y_pred,
        labels=classes,
        target_names=names,
        zero_division=0
    ))

    # 5. Confusion matrix
    cm = confusion_matrix(y_test, y_pred, labels=classes)
    print("CONFUSION MATRIX:\n", cm)

    # 6. Plot & save confusion matrix
    plt.figure()
    plt.imshow(cm, interpolation="nearest", cmap=plt.cm.Blues)
    plt.title("Confusion Matrix")
    plt.colorbar()
    plt.xticks(range(len(classes)), names)
    plt.yticks(range(len(classes)), names)
    for i in range(len(classes)):
        for j in range(len(classes)):
            plt.text(j, i, cm[i, j], ha='center', va='center')
    plt.xlabel("Predicted")
    plt.ylabel("True")
    plt.tight_layout()
    os.makedirs("results", exist_ok=True)
    plt.savefig("results/confusion_matrix.png")
    print("Saved confusion matrix to results/confusion_matrix.png\n")


def predict_file(filepath, model_path="src/models/best_model.h5", width=64):
    from src.preprocessing.convert_to_image import read_bytes, bytes_to_pixels

    # 1. Read & convert to pixels
    b = read_bytes(filepath)
    pix = bytes_to_pixels(b, width=width)

    # 2. Resize/pad height if needed to width×width
    h, w = pix.shape
    if h < width:
        pix = np.pad(pix, ((0, width-h), (0, 0)), 'constant', constant_values=0)
    elif h > width:
        pix = pix[:width, :]
    if w < width:
        pix = np.pad(pix, ((0, 0), (0, width-w)), 'constant', constant_values=0)
    elif w > width:
        pix = pix[:, :width]

    # 3. Normalize & reshape for model
    arr = pix.astype(np.float32) / 255.0
    arr = arr.reshape((1, width, width, 1))

    # 4. Load model & predict
    model = load_model(model_path)
    prob = model.predict(arr)[0, 0]
    label = "malware" if prob >= 0.5 else "benign"
    print(f"File: {filepath}\n→ Prediction: {label} (prob={prob:.4f})")


if __name__ == "__main__":
    if len(sys.argv) == 1:
        evaluate_test_set()
    else:
        for path in sys.argv[1:]:
            if os.path.isfile(path):
                predict_file(path)
            else:
                print(f"ERROR: File not found → {path}")