From 91437f30f5536e7f3ed0fe7978a47773f33730d6 Mon Sep 17 00:00:00 2001 From: Jason Date: Sun, 3 Apr 2022 11:58:31 -0700 Subject: [PATCH] prediction_dict --- load_estimation_model.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/load_estimation_model.py b/load_estimation_model.py index 51b8256..9891158 100644 --- a/load_estimation_model.py +++ b/load_estimation_model.py @@ -53,16 +53,17 @@ def load_estimation_model(inputfilename, outputfilename, begin, end, csv_export= model.load_state_dict(torch.load("em.pth")) my_prediction = model.forward(data) - my_prediction[0][0] = 1000 * float(my_prediction[0][0]) - my_prediction[0][1] = 1000 * float(my_prediction[0][1]) - my_prediction[0][2] = 1000 * float(my_prediction[0][2]) - my_prediction[0][3] = 1000 * float(my_prediction[0][3]) + prediction_dict = {} + prediction_dict["F1"] = 1000 * float(my_prediction[0][0]) + prediction_dict["F2"] = 1000 * float(my_prediction[0][1]) + prediction_dict["F3"] = 1000 * float(my_prediction[0][2]) + prediction_dict["F4"] = 1000 * float(my_prediction[0][3]) if csv_export: with open(outputfilename, "w") as wf: wf.write("NAME,begin,end,F1,F2,F3,F4\n") wf.write(name + "," + str(begin) + "," + str(end) + "," + \ - str(my_prediction[0][0]) + "," + str(my_prediction[0][1]) + "," + \ - str(my_prediction[0][2]) + "," + str(my_prediction[0][3]) + "\n") + str(prediction_dict["F1"]) + "," + str(prediction_dict["F2"]) + "," + \ + str(prediction_dict["F3"]) + "," + str(prediction_dict["F4"]) + "\n") - return my_prediction + return prediction_dict