Update my_model/state_manager.py

my_model/state_manager.py

@@ -24,10 +24,6 @@ class StateManager:
             st.session_state['button_label'] = "Load Model"
         if "previous_state" not in st.session_state:
             st.session_state['previous_state'] = {}
-        if "settings_changed" not in st.session_state:
-            st.session_state['settings_changed'] = self.settings_changed
-        if 'model_loaded' not in st.session_state:
-            st.session_state['model_loaded'] = False
         if 'loading_in_progress' not in st.session_state:
             st.session_state['loading_in_progress'] = False
         if 'load_button_clicked' not in st.session_state:
@@ -36,8 +32,12 @@ class StateManager:
             st.session_state['force_reload_button_clicked'] = False
         if 'time_taken_to_load_model' not in st.session_state:
             st.session_state['time_taken_to_load_model'] = 0
+        if "settings_changed" not in st.session_state:
+            st.session_state['settings_changed'] = self.settings_changed
+        if 'model_loaded' not in st.session_state:
+            st.session_state['model_loaded'] = self.is_model_loaded
 
-            
+    
 
 
     def set_up_widgets(self):
@@ -78,6 +78,7 @@ class StateManager:
         else:
             return col.slider(text, min_value, max_value, value, step, key=slider_key_name, disabled=self.is_widget_disabled)
 
+
     @property
     def is_widget_disabled(self):
         return st.session_state['loading_in_progress']
@@ -95,17 +96,16 @@ class StateManager:
         """
         return self.has_state_changed()
 
-    
+    @staticmethod
     def display_model_settings(self):
         """
         Displays a table of current model settings in the third column.
     
-        Uses formatted HTML to style the table for better readability.
         """
         self.col3.write("##### Current Model Settings:")
         data = [{'Setting': key, 'Value': str(value)} for key, value in st.session_state.items() if key in ["confidence_level", 'detection_model', 'method', 'kbvqa', 'previous_state', 'settings_changed', 'loading_in_progress', 'model_loaded', 'time_taken_to_load_model' ]]
         df = pd.DataFrame(data)
-        styled_df = df.style.set_properties(**{'background-color': 'white', 'color': 'black', 'border-color': 'black'}).set_table_styles([{'selector': 'th','props': [('background-color', 'gray'), ('font-weight', 'bold')]}])
+        #styled_df = df.style.set_properties(**{'background-color': 'white', 'color': 'black', 'border-color': 'black'}).set_table_styles([{'selector': 'th','props': [('background-color', 'gray'), ('font-weight', 'bold')]}])
         self.col3.write(df)
 
     
@@ -165,7 +165,7 @@ class StateManager:
         
         free_gpu_resources()
         
-        if self.is_model_loaded():
+        if self.is_model_loaded:
             try:
                 del st.session_state['kbvqa']
                 free_gpu_resources()
@@ -173,7 +173,6 @@ class StateManager:
                 free_gpu_resources()
                 pass
             
-
         
     # Function to check if any session state values have changed
     def has_state_changed(self):
@@ -197,7 +196,7 @@ class StateManager:
         """
         return st.session_state.get('kbvqa', None)
 
-    
+    @property
     def is_model_loaded(self):
         """
         Checks if the KBVQA model is loaded in the session state.
@@ -221,11 +220,13 @@ class StateManager:
         
         try:
             free_gpu_resources()
-            if self.is_model_loaded():
+            if self.is_model_loaded:
 
                 prepare_kbvqa_model(only_reload_detection_model=True)
                 st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level
                 self.col1.success("Model reloaded with updated settings and ready for inference.")
+                st.session_state['previous_state'] = {'method': st.session_state.method, 'detection_model': st.session_state.detection_model, 'confidence_level': st.session_state.confidence_level}
+                st.session_state['button_label'] = "Reload Model"
 
             free_gpu_resources()
         except Exception as e:
@@ -303,7 +304,26 @@ class StateManager:
             dict: The dictionary storing information about processed images.
         """
         return st.session_state['images_data']
+    
+    
+    def update_image_data(self, image_key, caption, detected_objects_str, analysis_done):
+        """
+        Updates the information stored for a specific image in the `images_data` dictionary in the application session state.
+    
+        Args:
+            image_key (str): Unique key for the image.
+            caption (str): The generated caption for the image.
+            detected_objects_str (str): String representation of detected objects.
+            analysis_done (bool): Flag indicating if analysis of the image is complete.
+        """
+        if image_key in st.session_state['images_data']:
+            st.session_state['images_data'][image_key].update({
+                'caption': caption,
+                'detected_objects_str': detected_objects_str,
+                'analysis_done': analysis_done
+            })
 
+    @staticmethod
     def resize_image(self, image_input, new_width=None, new_height=None):
         """
         Resize an image. If only new_width is provided, the height is adjusted to maintain aspect ratio.
@@ -346,20 +366,14 @@ class StateManager:
         return resized_image
 
     
-    
-    def update_image_data(self, image_key, caption, detected_objects_str, analysis_done):
-        """
-        Updates the information stored for a specific image in the `images_data` dictionary in the application session state.
-    
-        Args:
-            image_key (str): Unique key for the image.
-            caption (str): The generated caption for the image.
-            detected_objects_str (str): String representation of detected objects.
-            analysis_done (bool): Flag indicating if analysis of the image is complete.
-        """
-        if image_key in st.session_state['images_data']:
-            st.session_state['images_data'][image_key].update({
-                'caption': caption,
-                'detected_objects_str': detected_objects_str,
-                'analysis_done': analysis_done
-            })
+    @staticmethod
+    def display_message(self, message, message_type):
+        if message_type == "warning":
+            st.warning(message)
+        elif message_type == "text":
+            st.text(message)
+        elif message_type == "success":
+            st.success(messae)
+        elif message_type == "write":
+            st.write(message)
+        else: st.error("Message type unknown")