netsec-lab/competition/classifier.py
2021-06-05 15:18:06 +02:00

25 lines
633 B
Python

# 1. Importing new CSV data in pandas dataframes
import pandas as pd
import pickle
from sklearn import preprocessing
data = pd.read_csv("input.csv")
x = data.to_numpy()
# Preprocessing data - encode ip addresses to numerical values
le = preprocessing.LabelEncoder()
le.fit(data['sourceIPAddress'])
data['sourceIPAddress'] = le.transform(data['sourceIPAddress'])
le.fit(data['destinationIPAddress'])
data['destinationIPAddress'] = le.transform(data['destinationIPAddress'])
# 3. Loading a trained model
model = pickle.load(open('network_traffic_classifier.sav', 'rb'))
y_pred = model.predict(x)
print(data)
data.append(y_pred)