Update app.py

app.py

@@ -104,7 +104,11 @@ with row3_col5:
     st.image(img_8, use_column_width=True)
     if st.button(f"Select Image ↑", key="button_8"):
         st.session_state.image = img_8
-    
+
+# model
+if st.session_state.model is None:
+    session_state.model = YOLO("/home/user/app/models/head_model/best.pt")  # load        
+
 #####
 ###
 # Output and Interaction
@@ -134,9 +138,8 @@ if st.session_state.image is not None:
     st.write("\n\n")
     components.html("""<hr style="height:2px;border:none;color:#333;background-color:#333;" /> """)
 
-    # model
-    model = YOLO("/home/user/app/models/head_model/best.pt")  # load
-    results = model.predict(tiles) # predict
+    # predictions
+    results = session_state.model.predict(tiles) # predict
     
     
     _, row4_col2, row4_col3, _ = st.columns([1,4,4,1], gap="medium")
@@ -145,11 +148,11 @@ if st.session_state.image is not None:
     with row4_col3:
         row4a_col3, row4b_col3= st.columns([1,1], gap="medium")
         with row4a_col3:
-            if st.button("Back")
+            if st.button("Back"):
                 if session_state.image_counter > 0:
                     session_state.image_counter -= 1
         with row4b_col3:
-            if st.button("Next")
+            if st.button("Next"):
                 if session_state.image_counter < len(results):
                     session_state.image_counter -= 1