Update app.py

app.py

@@ -51,6 +51,7 @@ _, 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])
 _, row4_col2, row4_col3, _ = st.columns([1,4,4,1], gap="medium")
+row4a_col0, row4a_col1, row4a_col2= st.columns([1,1,1], gap="medium")
 
 
 
@@ -212,23 +213,25 @@ if st.session_state.screen == 2:
         st.image(st.session_state.image, use_column_width=True)
     
     with row4_col3:
-        row4a_col3, row4b_col3= st.columns([1,1], gap="medium")
-        with row4a_col3:
+        # navigation
+        with row4a_col0:
             if st.button("Back"):
                 if st.session_state.image_counter > 0:
                     st.session_state.image_counter -= 1
-        with row4b_col3:
+        with row4a_col1:
             if st.button("Select New Photo"):
                 st.session_state.screen = 1
                 st.session_state.results = None
                 st.rerun()
-        with row4b_col3:
+        with row4a_col2:
             if st.button("Next"):
-                if st.session_state.image_counter < len(st.session_state.results):
+                if st.session_state.image_counter < len(st.session_state.results - 1:
                     st.session_state.image_counter += 1
 
         # predictions
         if 'results' not in st.session_state:
+            st.write("The model is working. Please be patient...")
+            st.write("This process takes around 30 seconds.")
             st.session_state.results = st.session_state.model.predict(tiles, conf=0.5) # predict            
         
         # plot predicted image