network.py 3.2 KB
Newer Older
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
from pathlib import Path
from typing import List, Tuple

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle
from tensorflow import keras

from src.data.loader import np_arrays_from_scored_csv


class NeuralNetwork:
    def __init__(self, input_features: int = 2):
        self.model = keras.Sequential([
             keras.layers.Dense(10, activation='relu', input_shape=(input_features,)),
             keras.layers.Dense(12, activation='relu'),
             keras.layers.Dense(1, activation='sigmoid'),
            ]
        )
        self.model.compile(loss='binary_crossentropy',
                           optimizer='adam',
                           metrics=['accuracy'])
        self.false_positives = None

    def train(self,
              attacker_features_x: List[List[float]],
              benign_data: Tuple[np.ndarray, np.ndarray]):
        x, y = benign_data

        # There are some attacker's features
        attacker_features_x = np.array(attacker_features_x)
        if len(attacker_features_x[0]):
            attacker_features_y = [1 for _ in attacker_features_x]
            x = np.concatenate((x, attacker_features_x), axis=0)
            y = np.concatenate((y, attacker_features_y), axis=0)

        x, y = shuffle(x, y, random_state=1)
        self.model.fit(x, y, epochs=40, class_weight={0: 1, 1: 5})  # Train the model

    def calc_n0_false_positives(self, x_test: np.ndarray):
        res = self.model.predict(x_test)
        self.false_positives = sum(map(lambda x: 0 if x <= 0.5 else 1, res))

    def predict(self, attacker_features: List[float]) -> int:
        features = np.array([attacker_features])
        prediction = self.model.predict(features)
        # 1 -> malicious | 0 -> benign
        return 0 if prediction[0][0] <= 0.5 else 1

    def get_false_positive_rate(self):
        return self.false_positives

    def __str__(self):
        return f'(Neural network {self.__hash__()} with FP n0: {self.false_positives})'


if __name__ == '__main__':
    #  ---------------------------------------- Prepare the data ----------------------------------------
    # benign_x, benign_y = np_arrays_from_scored_csv(Path('../data/scored/all_benign_scored.csv'), 0, 1000)
    # malicious_x, malicious_y = np_arrays_from_scored_csv(Path('../data/scored/scored_malicious.csv'), 1, 50)

    benign_x, benign_y = np_arrays_from_scored_csv(Path('../data/scored/all_benign_scored.csv'), 0, 100)
    malicious_x, malicious_y = np_arrays_from_scored_csv(Path('../data/scored/scored_malicious.csv'), 1, 10)

    x = np.concatenate((benign_x, malicious_x), axis=0)
    y = np.concatenate((benign_y, malicious_y), axis=0)
    x, y = shuffle(x, y, random_state=1)

    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.10, random_state=41)

    # Initialize the model
    network = NeuralNetwork()
    network.model.fit(x_train, y_train, epochs=40)

    test_loss, test_acc = network.model.evaluate(x_test, y_test)
    print('Test accuracy:', test_acc)

    print(network.calc_n0_false_positives(x_test))
    # print(network.model.predict(np.array([[0.85, 0.98]]))[0])
    # Make a png model
    # keras.utils.plot_model(network.model, to_file='model.png', show_shapes=True)