File size: 1,154 Bytes
da8dab8 | 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 | import numpy as np
import torch
import joblib
from model.model import VOIPClassifier, get_top_features
scaler = joblib.load("scaler.pkl")
def load_model(model_path, input_dim):
model = VOIPClassifier(input_dim)
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()
return model
def predict(model, input_list, scaler):
input_array = np.array(input_list).reshape(1, -1)
input_scaled = scaler.transform(input_array)
input_tensor = torch.FloatTensor(input_scaled)
with torch.no_grad():
output = model(input_tensor)
probability = output.item()
predicted_class = 1 if probability >= 0.5 else 0
return {"probability": probability, "class": predicted_class}
def inference(input):
top_features = get_top_features()
model = load_model("VOIP_Classifier.pth", input_dim=len(top_features))
result = predict(model, input, scaler)
return result
if __name__ == "__main__":
top_features = get_top_features()
model = load_model("VOIP_Classifier.pth", input_dim=len(top_features))
sample_input = [2, 4, 5]
print(inference(sample_input))
|