Update app.py

app.py

@@ -54,13 +54,14 @@ if st.button("Forecast"):
         st.write(st.session_state.arima_forecast)
 
         # Plot ARIMA results
-        plt.figure(figsize=(10, 6))
-        plt.plot(st.session_state.arima_forecast, label='ARIMA Forecast', color='blue')
-        plt.title(f'ARIMA Forecast for Product {selected_product}')
-        plt.xlabel('Weeks')
-        plt.ylabel('Sales Quantity')
-        plt.legend()
-        st.pyplot()
+        fig, ax = plt.subplots(figsize=(10, 6))
+        ax.plot(st.session_state.arima_forecast, label='ARIMA Forecast', color='blue')
+        ax.set_title(f'ARIMA Forecast for Product {selected_product}')
+        ax.set_xlabel('Weeks')
+        ax.set_ylabel('Sales Quantity')
+        ax.legend()
+        
+        st.pyplot(fig)
 
     elif model_option == "Decision Tree":
         # Load Decision Tree model
@@ -76,13 +77,14 @@ if st.button("Forecast"):
         st.write(st.session_state.dt_forecast)
 
         # Plot Decision Tree results
-        plt.figure(figsize=(10, 6))
-        plt.plot(st.session_state.dt_forecast, label='Decision Tree Forecast', color='green')
-        plt.title(f'Decision Tree Forecast for Product {selected_product}')
-        plt.xlabel('Weeks')
-        plt.ylabel('Sales Quantity')
-        plt.legend()
-        st.pyplot()
+        fig, ax = plt.subplots(figsize=(10, 6))
+        ax.plot(st.session_state.dt_forecast, label='Decision Tree Forecast', color='green')
+        ax.set_title(f'Decision Tree Forecast for Product {selected_product}')
+        ax.set_xlabel('Weeks')
+        ax.set_ylabel('Sales Quantity')
+        ax.legend()
+        
+        st.pyplot(fig)
 
     elif model_option == "XGBoost":
         # Load XGBoost model
@@ -98,13 +100,14 @@ if st.button("Forecast"):
         st.write(st.session_state.xgb_forecast)
 
         # Plot XGBoost results
-        plt.figure(figsize=(10, 6))
-        plt.plot(st.session_state.xgb_forecast, label='XGBoost Forecast', color='orange')
-        plt.title(f'XGBoost Forecast for Product {selected_product}')
-        plt.xlabel('Weeks')
-        plt.ylabel('Sales Quantity')
-        plt.legend()
-        st.pyplot()
+        fig, ax = plt.subplots(figsize=(10, 6))
+        ax.plot(st.session_state.xgb_forecast, label='XGBoost Forecast', color='orange')
+        ax.set_title(f'XGBoost Forecast for Product {selected_product}')
+        ax.set_xlabel('Weeks')
+        ax.set_ylabel('Sales Quantity')
+        ax.legend()
+        
+        st.pyplot(fig)
 
 # Comparison button
 if st.button("Compare Models"):
@@ -136,23 +139,24 @@ if st.button("Compare Models"):
     st.session_state.xgb_forecast = xgb_model.predict(future_data)
 
     # Now plot all the forecasts
-    plt.figure(figsize=(10, 6))
+    fig, ax = plt.subplots(figsize=(10, 6))
 
     # ARIMA forecast
     if st.session_state.arima_forecast is not None:
-        plt.plot(st.session_state.arima_forecast, label='ARIMA Forecast', color='blue')
-
+        ax.plot(st.session_state.arima_forecast, label='ARIMA Forecast', color='blue')
+    
     # Decision Tree forecast
     if st.session_state.dt_forecast is not None:
-        plt.plot(st.session_state.dt_forecast, label='Decision Tree Forecast', color='green')
-
+        ax.plot(st.session_state.dt_forecast, label='Decision Tree Forecast', color='green')
+    
     # XGBoost forecast
     if st.session_state.xgb_forecast is not None:
-        plt.plot(st.session_state.xgb_forecast, label='XGBoost Forecast', color='orange')
-
-    plt.title(f'Comparison of Forecasts for Product {selected_product}')
-    plt.xlabel('Weeks')
-    plt.ylabel('Sales Quantity')
-    plt.legend()
-    st.pyplot()
+        ax.plot(st.session_state.xgb_forecast, label='XGBoost Forecast', color='orange')
+    
+    ax.set_title(f'Comparison of Forecasts for Product {selected_product}')
+    ax.set_xlabel('Weeks')
+    ax.set_ylabel('Sales Quantity')
+    ax.legend()
+    
+    st.pyplot(fig)