from gradio to streamlit

app.py

@@ -1,25 +1,18 @@
-import gradio as gr
+import streamlit as st
 import torch
 import numpy as np
 import matplotlib.pyplot as plt
+import seaborn as sns
 import io
-import base64
 import pickle
 import torch.nn as nn
 from sklearn.preprocessing import MinMaxScaler
-import pandas as pd
-
-df_weekly_km = pd.read_csv('Weekly_Km.csv')
-
 
 INPUT_SIZE = 1
 HIDDEN_LAYER_SIZE = 100
 OUTPUT_SIZE = 1
 
-data = df_weekly_km['Total Kilometers'].values.astype(float)
 scaler = MinMaxScaler(feature_range=(-1, 1))
-data_normalized = scaler.fit_transform(data.reshape(-1, 1))
-
 seq_length = 4
 
 class LSTM(nn.Module):
@@ -36,68 +29,54 @@ class LSTM(nn.Module):
         predictions = self.linear(lstm_out.view(len(input_seq), -1))
         return predictions[-1]
 
+with open('LSTM_MODEL.pkl', 'rb') as f:
+    model = pickle.load(f)
+
 def prepare_custom_input(last_values, seq_length, scaler):
     last_values_normalized = scaler.transform(np.array(last_values).reshape(-1, 1))
     input_seq = torch.from_numpy(last_values_normalized).float()
     return input_seq.view(-1)
 
-model_path = 'lstm_model.pth'
-
-model = LSTM(INPUT_SIZE, HIDDEN_LAYER_SIZE, OUTPUT_SIZE)
-
-model.load_state_dict(torch.load(model_path))
-
+model.eval()
 
-def predict_and_plot(week4, week3, week2, week1):
-    last_four_weeks = [week4, week3, week2, week1]
+def predict_and_plot(week1, week2, week3, week4):
+    last_four_weeks = [week1, week2, week3, week4]
 
     custom_input = prepare_custom_input(last_four_weeks, seq_length, scaler)
 
-    model.eval()
     with torch.no_grad():
         model.hidden_cell = (torch.zeros(1, 1, model.hidden_layer_size),
                              torch.zeros(1, 1, model.hidden_layer_size))
         prediction = model(custom_input)
 
     predicted_kilometers = scaler.inverse_transform(prediction.numpy().reshape(-1, 1))
-
-    predicted_value = f"{predicted_kilometers[0][0]:.2f}"
+    predicted_value = predicted_kilometers[0][0]
 
     weeks = ['Week -4', 'Week -3', 'Week -2', 'Week -1', 'Predicted Week']
-    actual_values = last_four_weeks + [predicted_value]
-    
-    plt.figure(figsize=(10, 6))
-    plt.plot(weeks, actual_values, marker='o', label='Total Kilometers')
+    actual_values = [week1, week2, week3, week4, predicted_value]
+
+    plt.figure(figsize=(8, 5))
+    sns.lineplot(x=weeks, y=actual_values, marker='o', label='Total Kilometers')
     plt.axvline(x='Predicted Week', color='red', linestyle='--', label='Predicted Week')
     plt.title('Total Kilometers for Last 4 Weeks and Prediction')
     plt.xlabel('Weeks')
     plt.ylabel('Total Kilometers')
-    plt.ylim(bottom=0)
-    plt.grid()
-    plt.tight_layout()
+    plt.grid(True)
     plt.legend()
 
-    # Saving the plot to a BytesIO object
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
-    plt.close()
-
-    # Converting to base64 for Gradio
-    img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-    img_tag = f'<img src="data:image/png;base64,{img_base64}"/>'
-    
-    return predicted_value, img_tag
-
-inputs = [gr.Number(label='Week -4'), 
-          gr.Number(label='Week -3'), 
-          gr.Number(label='Week -2'), 
-          gr.Number(label='Week -1')]
-outputs = [gr.Textbox(label='Predicted Week'), 
-           gr.HTML(label='Plot')]
-
-gr.Interface(fn=predict_and_plot, 
-             inputs=inputs, 
-             outputs=outputs,
-             title="LSTM Model Prediction",
-             description="Enter the total kilometers for the last 4 weeks to get the prediction for the next week along with a plot.").launch(share=True)
+    st.pyplot(plt)
+
+    return predicted_value
+
+# Streamlit interface
+st.title("LSTM Model Prediction")
+st.write("Enter the total kilometers for the last 4 weeks to get the prediction for the next week along with a plot.")
+
+week1 = st.number_input('Week -4', min_value=0)
+week2 = st.number_input('Week -3', min_value=0)
+week3 = st.number_input('Week -2', min_value=0)
+week4 = st.number_input('Week -1', min_value=0)
+
+if st.button('Predict'):
+    predicted_value = predict_and_plot(week1, week2, week3, week4)
+    st.write(f"Predicted Total Kilometers for the next week: {predicted_value:.2f}")