Update my_model/utilities/state_manager.py

my_model/utilities/state_manager.py

@@ -6,16 +6,20 @@ from my_model.KBVQA import KBVQA, prepare_kbvqa_model
 
 
 class StateManager:
-    #def __init__(self):
-       # self.initialize_state()
+    def __init__(self):
+        self.initialize_state()
 
     def initialize_state(self):
         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'] = {'selected_method': None, 'detection_model': None, 'confidence_level': None}
+            st.session_state['method'] = None
+        if 'detection_model': not in st.session_state
+            st.session_state['detection_model'] = None
         if 'kbvqa' not in st.session_state:
             st.session_state['kbvqa'] = None
+        if 'confidence_level': not in st.session_state
+            st.session_state['confidence_level'] = None
 
 
     def update_model_settings(self, detection_model=None, confidence_level=None, selected_method=None):