diff --git a/frontend/pages/prediction_regression.py b/frontend/pages/prediction_regression.py index 377274e..42acf34 100644 --- a/frontend/pages/prediction_regression.py +++ b/frontend/pages/prediction_regression.py @@ -1,6 +1,8 @@ import streamlit as st from sklearn.linear_model import LinearRegression import pandas as pd +import matplotlib.pyplot as plt +import numpy as np st.header("Prediction: Regression") @@ -25,5 +27,34 @@ if "data" in st.session_state: prediction = model.predict(pd.DataFrame([pred_values], columns=data_name)) st.write("Prediction:", prediction[0]) + + fig = plt.figure() + dataframe_sorted = pd.concat([X, y], axis=1).sort_values(by=data_name) + + if len(data_name) == 1: + X = dataframe_sorted[data_name[0]] + y = dataframe_sorted[target_name] + + prediction_array_y = [ + model.predict(pd.DataFrame([[dataframe_sorted[data_name[0]].iloc[i]]], columns=data_name))[0] + for i in range(dataframe_sorted.shape[0]) + ] + + plt.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[target_name], color='b') + plt.scatter(dataframe_sorted[data_name[0]], prediction_array_y, color='r') + else: + ax = fig.add_subplot(111, projection='3d') + + prediction_array_y = [ + model.predict(pd.DataFrame([[dataframe_sorted[data_name[0]].iloc[i], dataframe_sorted[data_name[1]].iloc[i]]], columns=data_name))[0] + for i in range(dataframe_sorted.shape[0]) + ] + + ax.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[data_name[1]], dataframe_sorted[target_name], color='b') + ax.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[data_name[1]], prediction_array_y, color='r') + + st.pyplot(fig) + + else: st.error("File not loaded")