55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
import pandas as pd
|
|
import pickle
|
|
from tensorflow.keras.models import load_model
|
|
from Feature import extract_url_features
|
|
from collections import Counter
|
|
from scipy.stats import entropy
|
|
import tensorflow as tf
|
|
|
|
# 🔹 URL 엔트로피 계산 함수
|
|
def calculate_url_entropy(url):
|
|
counter = Counter(url)
|
|
probabilities = [count / len(url) for count in counter.values()]
|
|
return entropy(probabilities, base=2)
|
|
|
|
# 🔹 스케일러 불러오기
|
|
with open("scaler.pkl", "rb") as f:
|
|
scaler = pickle.load(f)
|
|
|
|
# 🔹 모델 불러오기
|
|
model = load_model("best_model.h5")
|
|
|
|
# 🔹 예측 함수
|
|
@tf.function(reduce_retracing=True)
|
|
def predict_with_model(model, input_data):
|
|
return model(input_data)
|
|
|
|
# 🔹 입력 URL 받기
|
|
url = input("URL입력 : ")
|
|
|
|
# 🔹 Feature.py에서 피처 추출
|
|
features = extract_url_features(url)
|
|
|
|
# 🔹 누락된 피처 보완
|
|
features['url_length'] = len(url)
|
|
features['url_entropy'] = calculate_url_entropy(url)
|
|
|
|
# 🔹 데이터프레임 생성 및 정렬
|
|
input_df = pd.DataFrame([features])
|
|
expected_columns = list(scaler.feature_names_in_)
|
|
input_df = input_df[expected_columns]
|
|
|
|
# 🔹 스케일링
|
|
input_scaled = scaler.transform(input_df)
|
|
|
|
# 🔹 예측
|
|
prediction = predict_with_model(model, input_scaled)
|
|
score = float(prediction.numpy()[0][0]) # 🔥 정확히 float으로 변환
|
|
|
|
# 🔹 출력
|
|
threshold = 0.5
|
|
if score > threshold:
|
|
print(f"악성 (악성일 확률: {score:.4f})")
|
|
else:
|
|
print(f"정상 (정상일 확률: {1 - score:.4f})")
|