model.py 786 B

1234567891011121314151617181920212223242526272829303132
  1. import pickle
  2. import pandas as pd
  3. # 加载模型
  4. def load_model():
  5. with open('model_optimize/pkl/RF_filt.pkl', 'rb') as f:
  6. return pickle.load(f)
  7. # 初始化模型
  8. model = load_model() # 模型会在model.py被导入时加载一次
  9. # 预测函数
  10. def predict(input_data: pd.DataFrame):
  11. # 假设输入数据已经准备好
  12. predictions = model.predict(input_data)
  13. return predictions.tolist()
  14. if __name__ == '__main__':
  15. # 测试 predict 函数
  16. x = pd.DataFrame([{
  17. "organic_matter": 5.2,
  18. "chloride": 3.1,
  19. "cec": 25.6,
  20. "h_concentration": 0.5,
  21. "hn": 12.4,
  22. "al_concentration": 0.8,
  23. "free_alumina": 1.2,
  24. "free_iron": 0.9,
  25. "delta_ph": -0.2
  26. }])
  27. print(predict(x)) # 预测结果