Update app.py

app.py

@@ -2,7 +2,9 @@
 ###
 # 00 Imports
 # 01 Setup
-# 02 Screen 0
+# 02 Screen 0 - Introduction
+# 03 Screen 1 - Upload or Select Image
+
 ###
 #####
 
@@ -45,13 +47,15 @@ if 'model' not in st.session_state:
 
 _, row0_col1, _ = st.columns([2,3,2])
 _, row1_col1, _ = st.columns([2,3,2])
+_, row2_col1, row2_col2, row2_col3, row2_col4, _ = st.columns([8,3,3,3,3,8])
+_, row3_col1, row3_col2, row3_col3, row3_col4, _ = st.columns([8,3,3,3,3,8])
 
 
 
 
 #####
 ###
-# 02 Screen 0
+# 02 Screen 0 - Introduction
 ###
 #####
 # This displays to inform the user about the tool
@@ -60,29 +64,34 @@ _, row1_col1, _ = st.columns([2,3,2])
 # load model
 if st.session_state.screen == 0:
     with row0_col1:
+        # heading
         st.markdown("<h1 style='text-align: center; color: black;'>Find Waldo and His Friends</h1>", unsafe_allow_html=True)
         st.markdown("<h3 style='text-align: center; color: black;'>Using Computer Vision</h3>", unsafe_allow_html=True)
         st.markdown("<p style='text-align: left; color: black;'>This tool allows you to take a photo with your phone of a <i>Where's Waldo?</i> book page. Using the computer vision model YOLOv8-large, this tool find Waldo and his friends Wenda, Odlaw, Wizard, and Woof.</p>", unsafe_allow_html=True)
+
+        # proceed button
         if st.button(f"Proceed", key="Proceed"):
             st.session_state.screen = 1 
             st.rerun()
+        
+        # image of Waldo and his friends
         st.image("/home/user/app/home_page_images/Waldo_Friends.jpg", use_column_width=True)
        
     
     
 #####
 ###
-# Upload
+# 03 Screen 1 - Upload or Select Image
 ###
 #####
-if st.session_state.screen == 1:   
+if st.session_state.screen == 1:  
+    # heading and instructions
     with row0_col1:
         st.markdown("<h1 style='text-align: center; color: black;'>Find Waldo and His Friends</h1>", unsafe_allow_html=True)
         st.markdown("<h3 style='text-align: center; color: black;'>Using Computer Vision</h3>", unsafe_allow_html=True)
-        st.markdown("<p style='text-align: center; color: black;'>Upload an image using your cell phone (single book page), or select an imagefrom the below list</p>", unsafe_allow_html=True)
-
-
+        st.markdown("<p style='text-align: center; color: black;'>Upload an image using your cell phone (single book page), or select an image from the below list</p>", unsafe_allow_html=True)
 
+    # image uploader
     with row1_col1:
         uploaded_image = st.file_uploader(label="", 
                                           type=['jpg'], 
@@ -100,60 +109,66 @@ if st.session_state.screen == 1:
             uploaded_image = ImageOps.exif_transpose(uploaded_image)
             st.session_state.image = uploaded_image
 
+    # image selector
     st.markdown("<h4 style='text-align: center; color: black;'>Or Select One of the Following</h4>", unsafe_allow_html=True)
     
-    _, row2_col2, row2_col3, row2_col4, row2_col5, _ = st.columns([8,3,3,3,3,8])
-    _, row3_col2, row3_col3, row3_col4, row3_col5, _ = st.columns([8,3,3,3,3,8])
-    
     image_path = "/home/user/app/test_images/"
-    with row2_col2:
+    with row2_col1:
         img_1 = Image.open(image_path + "IMG_5356.JPG")
         img_1 = ImageOps.exif_transpose(img_1)
         st.image(img_1, use_column_width=True)
         if st.button(f"Select Image ↑", key="button_1"):
             st.session_state.image = img_1
-    with row2_col3:
+            st.session_state.screen = 2
+    with row2_col2:
         img_2 = Image.open(image_path + "IMG_5357.JPG")
         img_2 = ImageOps.exif_transpose(img_2)
         st.image(img_2, use_column_width=True)
         if st.button(f"Select Image ↑", key="button_2"):
             st.session_state.image = img_2
-    with row2_col4:
+            st.session_state.screen = 2
+    with row2_col3:
         img_3 = Image.open(image_path + "IMG_5368.JPG")
         img_3 = ImageOps.exif_transpose(img_3)
         st.image(img_3, use_column_width=True)
         if st.button(f"Select Image ↑", key="button_3"):
             st.session_state.image = img_3
-    with row2_col5:
+            st.session_state.screen = 2
+    with row2_col4:
         img_4 = Image.open(image_path + "IMG_5369.JPG")
         img_4 = ImageOps.exif_transpose(img_4)
         st.image(img_4, use_column_width=True)
         if st.button(f"Select Image ↑", key="button_4"):
             st.session_state.image = img_4
-    with row3_col2:
+            st.session_state.screen = 2
+    with row3_col1:
         img_5 = Image.open(image_path + "IMG_5382.JPG")
         img_5 = ImageOps.exif_transpose(img_5)
         st.image(img_5, use_column_width=True)
         if st.button(f"Select Image ↑", key="button_5"):
             st.session_state.image = img_5
-    with row3_col3:
+            st.session_state.screen = 2
+    with row3_col2:
         img_6 = Image.open(image_path + "IMG_5383.JPG")
         img_6 = ImageOps.exif_transpose(img_6)
         st.image(img_6, use_column_width=True)
         if st.button(f"Select Image ↑", key="button_6"):
             st.session_state.image = img_6
-    with row3_col4:
+            st.session_state.screen = 2
+    with row3_col3:
         img_7 = Image.open(image_path + "IMG_5408.JPG")
         img_7 = ImageOps.exif_transpose(img_7)
         st.image(img_7, use_column_width=True)
         if st.button(f"Select Image ↑", key="button_7"):
             st.session_state.image = img_7
-    with row3_col5:
+            st.session_state.screen = 2
+    with row3_col4:
         img_8 = Image.open(image_path + "IMG_5409.JPG")
         img_8 = ImageOps.exif_transpose(img_8)
         st.image(img_8, use_column_width=True)
         if st.button(f"Select Image ↑", key="button_8"):
-            st.session_state.image = img_8   
+            st.session_state.image = img_8 
+            st.session_state.screen = 2
 
 #####
 ###