Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| 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__() | |
| self.hidden_layer_size = hidden_layer_size | |
| self.lstm = nn.LSTM(input_size, hidden_layer_size) | |
| self.linear = nn.Linear(hidden_layer_size, output_size) | |
| self.hidden_cell = (torch.zeros(1, 1, self.hidden_layer_size), | |
| torch.zeros(1, 1, self.hidden_layer_size)) | |
| def forward(self, input_seq): | |
| lstm_out, self.hidden_cell = self.lstm(input_seq.view(len(input_seq), 1, -1), self.hidden_cell) | |
| predictions = self.linear(lstm_out.view(len(input_seq), -1)) | |
| return predictions[-1] | |
| # 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) | |
| # 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)) | |
| # 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)) | |
| prediction = model(custom_input) | |
| predicted_kilometers = scaler.inverse_transform(prediction.numpy().reshape(-1, 1)) | |
| 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') | |
| 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.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 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) |