test.py

import os
import time
import numpy as np
import pandas as pd
import pickle
from multiprocessing import Manager, Pool

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf

with open("/home/jovyan/models/hash_list.pkl", "rb") as f:
    hash_list = pickle.load(f)

# ---------------------直方图------------------------

start_time = time.time()

print("Histogram predict: {0:.2f}s".format(time.time()-start_time))

test_path = [os.path.join("/home/jovyan/histogram", sp) for sp in hash_list]

test_num = len(test_path)

raw_feature = np.empty((test_num, 512))

for i, fp in enumerate(test_path):
    with open(fp+'.txt', 'r') as f:
        feature = f.readlines()
    feature = [float(his.strip()) for his in feature]
    raw_feature[i] = feature

model = tf.keras.models.load_model('/home/jovyan/models/histogram_0.97.h5')

histogram_test = model.predict(raw_feature)

# ---------------------PE静态特征------------------------

print("PE raw predict: {0:.2f}s".format(time.time()-start_time))
start_time = time.time()
with open("/home/jovyan/models/raw_feature.pkl", "rb") as fp:
    pe_raw_models = pickle.load(fp)

with open("/home/jovyan/models/raw_feature_names.pkl", "rb") as fp:
    raw_feature_names = pickle.load(fp)

with open("/home/jovyan/models/rfc_pe_model.pkl", "rb") as f:
    rfc_pe_model = pickle.load(f)

with open("/home/jovyan/pe_raw/pe_raw_vectors.pkl", "rb") as f:
    pe_raw_vectors = pickle.load(f)

n_splits = 5
pe_raw_vectors = np.array(pe_raw_vectors, dtype=np.float32)

oof_test_skf = Manager().list([0] * n_splits) # np.empty((n_splits, test_num))
def pe_raw_predict(ind, model):
    oof_test_skf[ind] = model.predict(pe_raw_vectors)

stacking_test = []

for name in raw_feature_names:
    pool = Pool(5)
    for i, model in enumerate(pe_raw_models[name]):
        pool.apply_async(func=pe_raw_predict, args=(i, model))
    pool.close()
    pool.join()
    oof_test = np.array(list(oof_test_skf), dtype=np.float32).mean(axis=0)
    stacking_test.append(oof_test.reshape(-1, 1))

stacking_test = np.hstack(stacking_test)
raw_feature_test = rfc_pe_model.predict(stacking_test).reshape(-1, 1)

# ---------------------特征工程------------------------

print("Feature Engineering predict: {0:.2f}s".format(time.time()-start_time))
start_time = time.time()
with open("/home/jovyan/feature_engineering/feature_engineering_features.pkl", 'rb') as f:
    feature_engineering_features = pickle.load(f)

with open("/home/jovyan/models/keys.pkl", 'rb') as f:
    keys = pickle.load(f)

with open("/home/jovyan/models/lgb_models.pkl", "rb") as fp:
    lgb_models = pickle.load(fp)

train_df = pd.DataFrame(feature_engineering_features, columns=keys)

n_splits = 5

oof_test_skf = np.empty((n_splits, test_num))

for i, model in enumerate(lgb_models):
    oof_test_skf[i, :] = model.predict(train_df, num_iteration=model.best_iteration)

feature_engineerin_test = oof_test_skf.mean(axis=0).reshape(-1, 1)

# ---------------------融合------------------------

print("Final predict: {0:.2f}s".format(time.time()-start_time))
start_time = time.time()
with open("/home/jovyan/models/lr_rfc.pkl", "rb") as f:
    lr_rfc = pickle.load(f)

test = np.hstack([feature_engineerin_test, histogram_test , raw_feature_test])

labels_lr = lr_rfc[0].predict_proba(test)
labels_rfc = lr_rfc[1].predict_proba(test)

test_labels = []

for x, y in zip(labels_lr, labels_rfc):
    if x[1]*0.6+ y[1]*0.4  < 0.5:
        test_labels.append(0)
    else:
        test_labels.append(1)

print("Found {0} black samples. {1:.2f}s".format(sum(test_labels), time.time()-start_time))

result = []
for pt, label in zip(hash_list, test_labels):
    result.append("{0}, {1}\n".format(pt, label))
with open("/home/jovyan/malware_final.txt", 'w') as f:
    f.write(''.join(result).strip())