Update my_model/utilities/state_manager.py

my_model/utilities/state_manager.py

@@ -13,11 +13,10 @@ class StateManager:
         if 'images_data' not in st.session_state:
             st.session_state['images_data'] = {}
         if 'model_settings' not in st.session_state:
-            st.session_state['model_settings'] = {'detection_model': None, 'confidence_level': None}
+            st.session_state['model_settings'] = {'selected_method': None, 'detection_model': None, 'confidence_level': None}
         if 'kbvqa' not in st.session_state:
             st.session_state['kbvqa'] = None
-        if 'selected_method' not in st.session_state:
-            st.session_state['selected_method'] = None
+
 
     def update_model_settings(self, detection_model=None, confidence_level=None, selected_method=None):
         if detection_model is not None:
@@ -25,12 +24,12 @@ class StateManager:
         if confidence_level is not None:
             st.session_state['model_settings']['confidence_level'] = confidence_level
         if selected_method is not None:
-            st.session_state['selected_method'] = selected_method
+            st.session_state['model_settings']['selected_method'] = selected_method
 
     def check_settings_changed(self, current_selected_method, current_detection_model, current_confidence_level):
         return (st.session_state['model_settings']['detection_model'] != current_detection_model or
                 st.session_state['model_settings']['confidence_level'] != current_confidence_level or
-                st.session_state['selected_method'] != current_selected_method)
+                st.session_state['model_settings']['selected_method'] != current_selected_method)
 
     def display_model_settings(self):
         st.write("### Current Model Settings:")