Update app.py

app.py

@@ -38,43 +38,22 @@ _S3_PATH_OUTPUT = "s3://gretel-image-synthetics-use2/data/{identifier}/{image_ty
 # Eventually may want to move everything into functions and have a
 # if __name__ == "main" setup instead of everything inline.
 
-if "key" not in st.session_state:
-    st.session_state["key"] = uuid.uuid4().hex
-
-if "captcha_bool" not in st.session_state:
-    st.session_state["captcha_bool"] = False
-    
-if "model_inputs" not in st.session_state:
-    st.session_state["model_inputs"] = None
+bool_session_variables = ['captcha_bool', 'view', 'train_view', 'login', 'user_auth_sess', 'face_images_uploaded', 'theme_images_uploaded']
+for var in bool_session_variables:
+    if var not in st.session_state:
+        st.session_state[var] = False
 
-if (
-    "s3_face_file_path" not in st.session_state
-    and "s3_theme_file_path" not in st.session_state
-):
-    st.session_state["s3_face_file_path"] = None
-    st.session_state["s3_theme_file_path"] = None
+obj_session_variables = ['s3_face_file_path', 's3_theme_file_path', 'captcha_response', 'user_email']
+for var in obj_session_variables:
+    if var not in st.session_state:
+        st.session_state[var] = None
 
-if "view" not in st.session_state:
-    st.session_state["view"] = False
-
-if "train_view" not in st.session_state:
-    st.session_state["train_view"] = False
-
-if "captcha_response" not in st.session_state:
-    st.session_state["captcha_response"] = None
+if "key" not in st.session_state:
+    st.session_state["key"] = uuid.uuid4().hex
 
 if "captcha" not in st.session_state:
     st.session_state["captcha"] = {}
 
-if "login" not in st.session_state:
-    st.session_state["login"] = None
-
-if "user_auth_sess" not in st.session_state:
-    st.session_state["user_auth_sess"] = False
-
-if "user_email" not in st.session_state:
-    st.session_state["email_provided"] = True
-    
 def callback():
     st.session_state["button_clicked"] = True
 
@@ -232,26 +211,26 @@ def verify_captcha(captcha_id, captcha_response):
     return {"error": "Failed to verify captcha"}
 
 def train_model(model_inputs):
-    # Use banana backend
     api_key = "03cdd72e-5c04-4207-bd6a-fd5712c1740e"
     model_key = "1a3b4ce5-164f-4efb-9f4a-c2ad3a930d0b"
-    st.markdown(str(model_inputs))
     _ = banana.run(api_key, model_key, model_inputs)
 
-if st.session_state["email_provided"]:
+#first page - user provides email
+if st.session_state["user_email"] is None:
     user_email_input = st.empty()
-    with user_email_input.form(key='user_auth'):
+    with user_email_input.form(key='user_email_provided'):
         text_input = st.text_input(label='Please Enter Your Email')
         submit_button = st.form_submit_button(label='Submit')
         if submit_button: 
             st.session_state["user_auth_sess"] = True
-            st.session_state["email_provided"] = False
-            send_email(text_input, str(st.session_state["key"]))
-
+            st.session_state["user_email"] = text_input
+            send_email(st.session_state["user_email"], str(st.session_state["key"]))
+            user_email_input.empty()
+    
+#second page - user inputs one time code
 if st.session_state["user_auth_sess"]:
     user_auth = st.empty()
-    user_email_input.empty()
-    with user_auth.form("one-code"):
+    with user_auth.form("one-time-code"):
         text_input = st.text_input(label='Please Input One Time Code')
         submit_button = st.form_submit_button(label='Submit')
         if submit_button:
@@ -260,12 +239,12 @@ if st.session_state["user_auth_sess"]:
             else:
                 st.markdown("Please Enter Correct Code!")
 
+#third page - user inputs face images
 if st.session_state["login"]:
     identifier = st.session_state["key"]
     user_auth.empty()
-    user_email_input.empty()
     face_images = st.empty()
-    with face_images.form("my_form"):
+    with face_images.form("face_images_upload"):
         uploaded_files = st.file_uploader(
             "Choose image files", accept_multiple_files=True, type=["png", "jpg", "jpeg"]
         )
@@ -276,7 +255,11 @@ if st.session_state["login"]:
                     identifier, uploaded_files, "face"
                 )
             st.success(f'Uploading {len(uploaded_files)} files done!')
+            st.session_state["face_images_uploaded"] = True
 
+#fourth page - user inputs theme images
+if st.session_state["face_images_uploaded"]:
+    face_images.empty()
     preset_theme_images = st.empty()
     with preset_theme_images.form("choose-preset-theme"):
         img = image_select(
@@ -310,6 +293,7 @@ if st.session_state["login"]:
                         "num_images": 50,
                     }
                 st.success("Success!")
+                st.session_state["theme_images_uploaded"] = True
         with col2:
             submitted_4 = st.form_submit_button(
                 "If none of the themes interest you, click here!"
@@ -341,19 +325,18 @@ if st.session_state["login"]:
                         "num_images": 50,
                     }
                 st.success('Done!')
+                st.session_state["theme_images_uploaded"] = True
 
+#fifth page - user enters captcha and trains model
+if st.session_state["theme_images_uploaded"]:
+    if st.session_state["view"]:
+        custom_theme_images.empty()
+    preset_theme_images.empty()
     train = st.empty()
     with train.form("training"):
-        col1, col3, col2 = st.columns(3)
-        with col1: 
-            email = st.text_input("Enter Email")
-        with col2:
-            submitted = st.form_submit_button("Train Model!")
+        submitted = st.form_submit_button("Train Model!")
         if submitted:
-            if not email:
-                st.markdown('Please input an email!')
-            else:
-                st.session_state["captcha_bool"] = True
+            st.session_state["captcha_bool"] = True
 
     if st.session_state["captcha_bool"]:
         captcha_form = st.empty()
@@ -380,7 +363,7 @@ if st.session_state["login"]:
                     display_captcha = False
                     with st.spinner("Model Fine Tuning..."):
                         st.session_state["model_inputs"]["identifier"] = st.session_state["key"]
-                        st.session_state["model_inputs"]["email"] = email
+                        st.session_state["model_inputs"]["email"] = st.session_state["user_email"]
                         s3_output_path = _S3_PATH_OUTPUT.format(identifier=st.session_state["key"], image_type="generated")
                         # The backend does not have s3 credentials, so generate
                         # presigned urls for the backend to use to write and read