26 lines
893 B
Python
26 lines
893 B
Python
import os
|
|
import joblib
|
|
import pandas as pd
|
|
from sklearn.impute import SimpleImputer
|
|
|
|
TEST_DATA_DIR = "datasets_test"
|
|
|
|
test_df_list = []
|
|
for file in os.listdir(TEST_DATA_DIR):
|
|
file_path = os.path.join(TEST_DATA_DIR, file)
|
|
data_test = pd.read_csv(file_path, delim_whitespace=True, skiprows=1,
|
|
names=["tbid", "tphys", "r", "vr", "vt", "ik1", "ik2", "sm1", "sm2", "a", "e",
|
|
"collapsed"])
|
|
X_test = data_test.iloc[:, 1:-1].values
|
|
|
|
imputer = SimpleImputer(strategy='mean')
|
|
X_imputed = imputer.fit_transform(X_test)
|
|
|
|
model_filename = 'trained_model.pkl'
|
|
model = joblib.load(model_filename)
|
|
|
|
predictions = model.predict(X_test)
|
|
data_test['prediction'] = predictions
|
|
|
|
save_path = os.path.join(TEST_DATA_DIR, file + '.pred')
|
|
data_test.to_csv(save_path, index=False, sep=' ') |