Update app.py

app.py

@@ -1,3 +1,4 @@
+
 # Import required libraries
 import yfinance as yf
 import pandas as pd
@@ -7,6 +8,8 @@ from sklearn.preprocessing import MinMaxScaler
 import matplotlib.pyplot as plt
 import gradio as gr
 from datetime import datetime
+from io import BytesIO
+import PIL.Image
 
 # Step 1: Fetch stock data
 def fetch_stock_data(ticker, start_date, end_date):
@@ -118,6 +121,9 @@ def plot_predictions(data, predicted_prices, scaler):
         data (pd.DataFrame): DataFrame containing historical stock data.
         predicted_prices (list): Predicted stock prices for future dates.
         scaler (MinMaxScaler): Scaler to inverse transform the predicted prices.
+
+    Returns:
+        PIL.Image: The image of the plot saved in memory.
     """
     last_60_days = data['Close'][-60:].values.reshape(-1, 1)
     predicted_prices = np.array(predicted_prices).reshape(-1, 1)
@@ -134,7 +140,17 @@ def plot_predictions(data, predicted_prices, scaler):
     plt.xlabel("Days")
     plt.ylabel("Stock Price")
     plt.legend()
-    plt.show()
+
+    # Save the plot to a bytes buffer
+    buf = BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    image = PIL.Image.open(buf)
+
+    # Clear the plot so it doesn’t overlap
+    plt.clf()
+
+    return image
 
 # Step 7: Gradio Interface Function
 def stock_prediction_app(ticker, start_date_str, end_date_str):
@@ -148,7 +164,7 @@ def stock_prediction_app(ticker, start_date_str, end_date_str):
         end_date_str (str): End date selected by the user (YYYY-MM-DD).
 
     Returns:
-        None (Displays a plot of historical and predicted stock prices).
+        PIL.Image: The plot showing historical and predicted stock prices.
     """
     # Convert date strings to datetime objects
     start_date = datetime.strptime(start_date_str, "%Y-%m-%d").date()
@@ -168,7 +184,9 @@ def stock_prediction_app(ticker, start_date_str, end_date_str):
     predicted_prices = predict_future(model, scaled_data)
     
     # Plot historical and predicted prices
-    plot_predictions(data, predicted_prices, scaler)
+    image = plot_predictions(data, predicted_prices, scaler)
+    
+    return image
 
 # Step 8: Gradio UI Setup
 tickers = ["AAPL", "GOOGL", "MSFT", "AMZN", "TSLA", "META", "NFLX", "NVDA", "BABA", "BA"]
@@ -181,7 +199,7 @@ ui = gr.Interface(
         gr.Textbox(label="Start Date (YYYY-MM-DD)"),   # Use textbox instead
         gr.Textbox(label="End Date (YYYY-MM-DD)")      # Use textbox instead
     ],
-    outputs="plot",
+    outputs=gr.Image(),  # Updated output to return an image
     title="Stock Prediction App",
     description="Select a stock ticker and date range to predict future prices."
 )