Update app.py

app.py

@@ -33,6 +33,14 @@ def display_reconstruction(index, autoencoder, encoder, x_test):
 # Main Streamlit app
 st.title("Autoencoder Training and Visualization")
 
+# Initialize or reuse session state
+if "autoencoder" not in st.session_state:
+    st.session_state.autoencoder = None
+if "encoder" not in st.session_state:
+    st.session_state.encoder = None
+if "x_test" not in st.session_state:
+    st.session_state.x_test = None
+
 # Button to trigger training
 if st.button("Train Autoencoder"):
     # Load and preprocess data
@@ -42,7 +50,10 @@ if st.button("Train Autoencoder"):
     x_test = x_test.astype('float32') / 255.0
 
     x_train = np.reshape(x_train, (-1, 784))  # Flatten to (None, 784)
-    x_test = np.reshape(x_test, (-1, 784))
+    x_test = np reshape(x_test, (-1, 784))
+
+    # Save test set in session state
+    st.session_state.x_test = x_test
 
     # Define autoencoder architecture
     input_img = keras.Input(shape=(784,))
@@ -62,12 +73,14 @@ if st.button("Train Autoencoder"):
     with st.spinner("Training in progress..."):
         autoencoder.fit(x_train, x_train, epochs=5, batch_size=128, validation_data=(x_test, x_test))
 
-    # Create encoder model
-    encoder = keras.Model(input_img, latent_vector)
+    # Save trained models to session state
+    st.session_state.autoencoder = autoencoder
+    st.session_state.encoder = keras.Model(input_img, latent_vector)
 
-    # Input for image index to display
+# Input for image index to display
+if st.session_state.autoencoder:
     test_index = st.number_input("Enter an index (0-9999) to view an image from the test set:", min_value=0, max_value=9999)
 
     # Button to display the reconstruction
     if st.button("Display Reconstruction"):
-        display_reconstruction(test_index, autoencoder, encoder, x_test)
+        display_reconstruction(test_index, st.session_state.autoencoder, st.session_state.encoder, st.session_state.x_test)