Update app.py

app.py

@@ -43,6 +43,9 @@ if 'image' not in st.session_state:
 if 'image_counter' not in st.session_state:
     st.session_state.image_counter = 0 
 
+if 'confidence' not in st.session_state:
+    st.session_state.confidence = 50
+
 if 'model' not in st.session_state:
     st.session_state.model = YOLO("/home/user/app/models/head_model/best.pt")  # load 
 
@@ -204,16 +207,17 @@ if st.session_state.screen == 2:
                     st.session_state.image_counter += 1
 
         # predictions
-        if 'results' not in st.session_state:
+        if 'results' not in st.session_state or st.session_state.results = None:
             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(st.session_state.tiles, conf=0.5) # predict     
+            st.session_state.results = st.session_state.model.predict(st.session_state.tiles, conf=st.session_state.confidence) # predict     
             st.rerun()
         
         # 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)
+        st.session_state.confidence = st.slider(label="Change Model Confidence Level and Rerun", min_value=1, max_value=99, value=50, step=1)