Update app.py

app.py

@@ -29,6 +29,7 @@ from ultralytics import YOLO
 # 01 Setup
 ###
 #####
+# This loads general code for use by the tool
 st.set_page_config(layout="wide")
 
 uploaded_image = None
@@ -60,9 +61,6 @@ _, row4_col2, row4_col3, _ = st.columns([1,4,4,1], gap="medium")
 ###
 #####
 # This displays to inform the user about the tool
-
-
-# load model
 if st.session_state.screen == 0:
     with row0_col1:
         # heading
@@ -212,13 +210,6 @@ if st.session_state.screen == 2:
 
     with row4_col2:
         st.image(st.session_state.image, use_column_width=True)
-
-    # predictions
-    if 'results' not in st.session_state:
-        st.session_state.results = st.session_state.model.predict(tiles, conf=0.5) # predict
-    
-    
-    
     
     with row4_col3:
         row4a_col3, row4b_col3= st.columns([1,1], gap="medium")
@@ -226,12 +217,21 @@ if st.session_state.screen == 2:
             if st.button("Back"):
                 if st.session_state.image_counter > 0:
                     st.session_state.image_counter -= 1
+        with row4b_col3:
+            if st.button("Select New Photo"):
+                st.session_state.screen = 1
+                st.session_state.results = None
+                st.rerun()
         with row4b_col3:
             if st.button("Next"):
                 if st.session_state.image_counter < len(st.session_state.results):
                     st.session_state.image_counter += 1
-                
-        # Plot predicted image
+
+        # predictions
+        if 'results' not in st.session_state:
+            st.session_state.results = st.session_state.model.predict(tiles, conf=0.5) # predict            
+        
+        # plot predicted image
         im_bgr = st.session_state.results[st.session_state.image_counter].plot()  # BGR-order numpy array
         im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image
         st.image(im_rgb)