Fixed error

app.py

@@ -1,20 +1,30 @@
-import streamlit as st
+import gradio as gr
 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
 
+# Load your dataset
+df_weekly_km = pd.read_csv('Weekly_Km.csv')
+
+# Model parameters
 INPUT_SIZE = 1
 HIDDEN_LAYER_SIZE = 100
 OUTPUT_SIZE = 1
 
+# Normalize the data
+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
 
+# Define LSTM model
 class LSTM(nn.Module):
     def __init__(self, input_size=INPUT_SIZE, hidden_layer_size=HIDDEN_LAYER_SIZE, output_size=OUTPUT_SIZE):
         super(LSTM, self).__init__()
@@ -29,21 +39,24 @@ 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)
-
+# Function to prepare input
 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.eval()
+# Load the trained model
+model_path = 'lstm_model.pth'
+model = LSTM(INPUT_SIZE, HIDDEN_LAYER_SIZE, OUTPUT_SIZE)
+model.load_state_dict(torch.load(model_path))
 
-def predict_and_plot(week1, week2, week3, week4):
-    last_four_weeks = [week1, week2, week3, week4]
+# Prediction and plot function
+def predict_and_plot(week4, week3, week2, week1):
+    last_four_weeks = [week4, week3, week2, week1]
 
     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))
@@ -53,30 +66,41 @@ def predict_and_plot(week1, week2, week3, week4):
     predicted_value = predicted_kilometers[0][0]
 
     weeks = ['Week -4', 'Week -3', 'Week -2', 'Week -1', 'Predicted Week']
-    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')
+    actual_values = last_four_weeks + [predicted_value]
+    
+    plt.figure(figsize=(10, 6))
+    plt.plot(weeks, 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.grid(True)
+    plt.ylim(bottom=0)
+    plt.grid()
+    plt.tight_layout()
     plt.legend()
 
-    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}")
+    # 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 f"{predicted_value:.2f}", img_tag
+
+# Gradio interface
+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)