Update my_model/state_manager.py

my_model/state_manager.py

@@ -1,5 +1,6 @@
 # This module contains the StateManager class.
-# The StateManager class is primarily designed to facilitate the Run Inference tool that allows users to load, run, and test the models.
+# The StateManager class is primarily designed to facilitate the Run Inference tool that allows users to load, run, 
+# and test the models.
 
 
 import pandas as pd
@@ -14,14 +15,18 @@ from my_model.KBVQA import KBVQA, prepare_kbvqa_model
 
 class StateManager:
     """
-    Manages the user interface and session state for the Streamlit-based Knowledge-Based Visual Question Answering (KBVQA) application.
+    Manages the user interface and session state for the Streamlit-based Knowledge-Based Visual Question Answering 
+    (KBVQA) application.
 
-    This class includes methods to initialize the session state, set up various UI widgets for model selection and settings, 
+    This class includes methods to initialize the session state, set up various UI widgets for model selection and 
+    settings, 
     manage the loading and reloading of the KBVQA model, and handle the processing and analysis of images. 
     It tracks changes to the application's state to ensure the correct configuration is maintained. 
-    Additionally, it provides methods to display the current model settings and the complete application state within the Streamlit interface.
+    Additionally, it provides methods to display the current model settings and the complete application state within 
+    the Streamlit interface.
 
-    The StateManager class is primarily designed to facilitate the Run Inference tool that allows users to load, run, and test the models.
+    The StateManager class is primarily designed to facilitate the Run Inference tool that allows users to load, run, 
+    and test the models.
 
     Attributes:
         col1 (streamlit.columns): The first column in the Streamlit layout.
@@ -29,28 +34,26 @@ class StateManager:
         col3 (streamlit.columns): The third column in the Streamlit layout.
     """
 
-    
     def __init__(self) -> None:
         """
         Initializes the StateManager instance, setting up the Streamlit columns for the user interface.
         """
-        
+
         # Create three columns with different widths
-        self.col1, self.col2, self.col3 = st.columns([0.2, 0.6, 0.2])  
+        self.col1, self.col2, self.col3 = st.columns([0.2, 0.6, 0.2])
 
-    
     def initialize_state(self) -> None:
         """
         Initializes the Streamlit session state with default values for various keys.
         """
-        
+
         if "previous_state" not in st.session_state:
             st.session_state['previous_state'] = {'method': None, 'detection_model': None, 'confidence_level': None}
         if 'images_data' not in st.session_state:
             st.session_state['images_data'] = {}
         if 'kbvqa' not in st.session_state:
             st.session_state['kbvqa'] = None
-        if "button_label" not in st.session_state: 
+        if "button_label" not in st.session_state:
             st.session_state['button_label'] = "Load Model"
         if 'loading_in_progress' not in st.session_state:
             st.session_state['loading_in_progress'] = False
@@ -65,7 +68,6 @@ class StateManager:
         if 'model_loaded' not in st.session_state:
             st.session_state['model_loaded'] = self.is_model_loaded
 
-    
     def set_up_widgets(self) -> None:
         """
         Sets up user interface widgets for selecting models, settings, and displaying model settings conditionally.
@@ -74,18 +76,22 @@ class StateManager:
             None
         """
 
-        self.col1.selectbox("Choose a model:", ["13b-Fine-Tuned Model", "7b-Fine-Tuned Model", "Vision-Language Embeddings Alignment"], index=1, key='method', disabled=self.is_widget_disabled)
-        detection_model = self.col1.selectbox("Choose a model for objects detection:", ["yolov5", "detic"], index=1, key='detection_model', disabled=self.is_widget_disabled)
+        self.col1.selectbox("Choose a model:",
+                            ["13b-Fine-Tuned Model", "7b-Fine-Tuned Model", "Vision-Language Embeddings Alignment"],
+                            index=1, key='method', disabled=self.is_widget_disabled)
+        detection_model = self.col1.selectbox("Choose a model for objects detection:", ["yolov5", "detic"], index=1,
+                                              key='detection_model', disabled=self.is_widget_disabled)
         default_confidence = 0.2 if st.session_state.detection_model == "yolov5" else 0.4
-        self.set_slider_value(text="Select minimum detection confidence level", min_value=0.1, max_value=0.9, value=default_confidence, step=0.05, slider_key_name='confidence_level', col=self.col1)
+        self.set_slider_value(text="Select minimum detection confidence level", min_value=0.1, max_value=0.9,
+                              value=default_confidence, step=0.05, slider_key_name='confidence_level', col=self.col1)
 
         # Conditional display of model settings
-        show_model_settings = self.col3.checkbox("Show Model Settings", True,  disabled=self.is_widget_disabled)
+        show_model_settings = self.col3.checkbox("Show Model Settings", True, disabled=self.is_widget_disabled)
         if show_model_settings:
             self.display_model_settings
 
-    
-    def set_slider_value(self, text: str, min_value: float, max_value: float, value: float, step: float, slider_key_name: str, col=None) -> None:
+    def set_slider_value(self, text: str, min_value: float, max_value: float, value: float, step: float,
+                         slider_key_name: str, col=None) -> None:
         """
         Creates a slider widget with the specified parameters, optionally placing it in a specific column.
 
@@ -101,47 +107,46 @@ class StateManager:
         Returns:
             None
         """
-        
+
         if col is None:
-            return st.slider(text, min_value, max_value, value, step, key=slider_key_name, disabled=self.is_widget_disabledd)
+            return st.slider(text, min_value, max_value, value, step, key=slider_key_name,
+                             disabled=self.is_widget_disabledd)
         else:
-            return col.slider(text, min_value, max_value, value, step, key=slider_key_name, disabled=self.is_widget_disabled)
+            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) -> bool:
         """
         Checks if widgets should be disabled based on the 'loading_in_progress' state.
-        
+
         Returns:
             bool: True if widgets should be disabled, False otherwise.
         """
-        
+
         return st.session_state['loading_in_progress']
 
     def disable_widgets(self) -> None:
         """
         Disables widgets by setting the 'loading_in_progress' state to True.
-        
+
         Returns:
             None
         """
-        
+
         st.session_state['loading_in_progress'] = True
 
-        
     @property
     def settings_changed(self) -> bool:
         """
         Checks if any model settings have changed compared to the previous state.
-    
+
         Returns:
             bool: True if any setting has changed, False otherwise.
         """
-        
+
         return self.has_state_changed()
 
-        
     @property
     def confidance_change(self) -> bool:
         """
@@ -150,10 +155,9 @@ class StateManager:
         Returns:
             bool: True if the confidence level has changed, False otherwise.
         """
-        
+
         return st.session_state["confidence_level"] != st.session_state["previous_state"]["confidence_level"]
 
-        
     def update_prev_state(self) -> None:
         """
         Updates the 'previous_state' in the session state with the current state values.
@@ -161,15 +165,14 @@ class StateManager:
         Returns:
             None
         """
-        
+
         for key in st.session_state['previous_state']:
             st.session_state['previous_state'][key] = st.session_state[key]
 
-            
     def load_model(self) -> None:
         """
         Loads the KBVQA model based on the chosen method and settings.
-    
+
         - Frees GPU resources before loading.
         - Calls `prepare_kbvqa_model` to create the model.
         - Sets the detection confidence level on the model object.
@@ -179,7 +182,7 @@ class StateManager:
         Returns:
             None
         """
-        
+
         try:
             free_gpu_resources()
             st.session_state['kbvqa'] = prepare_kbvqa_model()
@@ -190,14 +193,14 @@ class StateManager:
             st.session_state['button_label'] = "Reload Model"
             free_gpu_resources()
             free_gpu_resources()
-            
+
         except Exception as e:
             st.error(f"Error loading model: {e}")
-            
 
     def force_reload_model(self) -> None:
         """
-        Forces a reload of all models, freeing up GPU resources. This method deletes the current models and calls `free_gpu_resources`.
+        Forces a reload of all models, freeing up GPU resources. This method deletes the current models and calls 
+        `free_gpu_resources`.
 
         - Deletes the current KBVQA model from the session state.
         - Calls `prepare_kbvqa_model` with `force_reload=True` to reload the model.
@@ -207,7 +210,7 @@ class StateManager:
         Returns:
             None
         """
-        
+
         try:
             self.delete_model()
             free_gpu_resources()
@@ -215,14 +218,13 @@ class StateManager:
             st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level
             # Update the previous state with current session state values
             self.update_prev_state()
-            
+
             st.session_state['model_loaded'] = True
             free_gpu_resources()
         except Exception as e:
             st.error(f"Error reloading model: {e}")
             free_gpu_resources()
 
-        
     def delete_model(self) -> None:
         """
         This method deletes the current models and calls `free_gpu_resources`.
@@ -230,9 +232,9 @@ class StateManager:
         Returns:
             None
         """
-        
+
         free_gpu_resources()
-        
+
         if self.is_model_loaded:
             try:
                 del st.session_state['kbvqa']
@@ -242,12 +244,11 @@ class StateManager:
                 free_gpu_resources()
                 free_gpu_resources()
                 pass
-            
-        
+
     def has_state_changed(self) -> bool:
         """
         Compares current session state with the previous state to identify changes.
-        
+
         Returns:
             bool: True if any change is found, False otherwise.
         """
@@ -255,41 +256,38 @@ class StateManager:
             if key == 'confidence_level':
                 continue  # confidence_level tracker is separate
             if key in st.session_state and st.session_state[key] != st.session_state['previous_state'][key]:
-                
                 return True  # Found a change
-        else: return False  # No changes found   
+        else:
+            return False  # No changes found   
 
-    
     def get_model(self) -> KBVQA.KBVQA():
         """
         Retrieves the KBVQA model from the session state.
-        
+
         Returns:
             KBVQA: The loaded KBVQA model, or None if not loaded.
         """
-        
+
         return st.session_state.get('kbvqa', None)
 
-                
     @property
     def is_model_loaded(self) -> bool:
         """
         Checks if the KBVQA model is loaded in the session state.
-    
+
         Returns:
             bool: True if the model is loaded, False otherwise.
         """
-        
+
         return 'kbvqa' in st.session_state and st.session_state['kbvqa'] is not None and \
-           st.session_state.kbvqa.all_models_loaded \
-           and (st.session_state['previous_state']['method'] is not None 
-                and st.session_state['method'] == st.session_state['previous_state']['method'])
+               st.session_state.kbvqa.all_models_loaded \
+               and (st.session_state['previous_state']['method'] is not None
+                    and st.session_state['method'] == st.session_state['previous_state']['method'])
 
-    
     def reload_detection_model(self) -> None:
         """
         Reloads only the detection model of the KBVQA model with updated settings.
-    
+
         - Frees GPU resources before reloading.
         - Checks if the model is already loaded.
         - Calls `prepare_kbvqa_model` with `only_reload_detection_model=True`.
@@ -299,7 +297,7 @@ class StateManager:
         Returns:
             None
         """
-        
+
         try:
             free_gpu_resources()
             if self.is_model_loaded:
@@ -313,18 +311,18 @@ class StateManager:
         except Exception as e:
             st.error(f"Error reloading detection model: {e}")
 
-    
     def process_new_image(self, image_key: str, image) -> None:
         """
-        Processes a new uploaded image by creating an entry in the `images_data` dictionary in the application session state.
-    
+        Processes a new uploaded image by creating an entry in the `images_data` dictionary in the application session 
+        state.
+
         This dictionary stores information about each processed image, including:
             - `image`: The original image data.
             - `caption`: Generated caption for the image.
             - `detected_objects_str`: String representation of detected objects.
             - `qa_history`: List of questions and answers related to the image.
             - `analysis_done`: Flag indicating if analysis is complete.
-    
+
         Args:
             image_key (str): Unique key for the image.
             image (obj): The uploaded image data.
@@ -332,7 +330,7 @@ class StateManager:
         Returns:
             None
         """
-        
+
         if image_key not in st.session_state['images_data']:
             st.session_state['images_data'][image_key] = {
                 'image': image,
@@ -342,8 +340,6 @@ class StateManager:
                 'analysis_done': False
             }
 
-
-    
     def analyze_image(self, image) -> Tuple[str, str, object]:
         """
         Analyzes the image using the KBVQA model.
@@ -352,7 +348,7 @@ class StateManager:
         - Displays a "Analyzing the image .." message.
         - Calls KBVQA methods to generate a caption and detect objects.
         - Returns the generated caption, detected objects string, and image with bounding boxes.
-    
+
         Args:
             image (obj): The image data to analyze.
 
@@ -368,11 +364,10 @@ class StateManager:
         free_gpu_resources()
         return caption, detected_objects_str, image_with_boxes
 
-    
     def add_to_qa_history(self, image_key: str, question: str, answer: str, prompt_length: int) -> None:
         """
         Adds a question-answer pair to the QA history of a specific image, to be used as a history tracker.
-    
+
         Args:
             image_key (str): Unique key for the image.
             question (str): The question asked about the image.
@@ -382,26 +377,25 @@ class StateManager:
         Returns:
             None
         """
-        
+
         if image_key in st.session_state['images_data']:
             st.session_state['images_data'][image_key]['qa_history'].append((question, answer, prompt_length))
 
-    
     def get_images_data(self) -> Dict:
         """
         Returns the dictionary containing processed image data from the session state.
-    
+
         Returns:
             dict: The dictionary storing information about processed images.
         """
-        
+
         return st.session_state['images_data']
-    
-    
+
     def update_image_data(self, image_key: str, caption: str, detected_objects_str: str, analysis_done: bool) -> None:
         """
-        Updates the information stored for a specific image in the `images_data` dictionary in the application session state.
-    
+        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.
@@ -418,21 +412,20 @@ class StateManager:
                 'analysis_done': analysis_done
             })
 
- 
     def resize_image(self, image_input, new_width: Optional[int] = None, new_height: Optional[int] = None) -> Image:
         """
         Resizes an image. If only new_width is provided, the height is adjusted to maintain aspect ratio.
         If both new_width and new_height are provided, the image is resized to those dimensions.
-    
+
         Args:
             image_input (PIL.Image.Image): The image to resize.
             new_width (int, optional): The target width of the image.
             new_height (int, optional): The target height of the image.
-    
+
         Returns:
             PIL.Image.Image: The resized image.
         """
-    
+
         img = copy.deepcopy(image_input)
         if isinstance(img, str):
             # Open the image from a file path
@@ -442,7 +435,7 @@ class StateManager:
             image = img
         else:
             raise ValueError("image_input must be a file path or a PIL Image object")
-                
+
         if new_width is not None and new_height is None:
             # Calculate new height to maintain aspect ratio
             original_width, original_height = image.size
@@ -455,17 +448,15 @@ class StateManager:
             new_width = int(original_width * ratio)
         elif new_width is None and new_height is None:
             raise ValueError("At least one of new_width or new_height must be provided")
-    
+
         # Resize the image
         resized_image = image.resize((new_width, new_height))
         return resized_image
 
-    
- 
     def display_message(self, message: str, message_type: str) -> None:
         """
         Displays a message in the Streamlit interface based on the specified message type.
-    
+
         Args:
             message (str): The message to display.
             message_type (str): The type of message ('warning', 'text', 'success', 'write', or 'error').
@@ -473,18 +464,18 @@ class StateManager:
         Returns:
             None
         """
-        
+
         if message_type == "warning":
             st.warning(message)
         elif message_type == "text":
             st.text(message)
         elif message_type == "success":
-            st.success(messae)
+            st.success(message)
         elif message_type == "write":
             st.write(message)
-        else: st.error("Message type unknown") 
-        
-        
+        else:
+            st.error("Message type unknown")
+
     @property
     def display_model_settings(self) -> None:
         """
@@ -494,22 +485,23 @@ class StateManager:
             None
         """
         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', 'images_data' ]]
+        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', 'images_data']]
         df = pd.DataFrame(data).reset_index(drop=True)
         return self.col3.write(df)
 
-    
     def display_session_state(self, col) -> None:
         """
         Displays a table of the complete application state in the specified column.
-    
+
         Args:
             col (streamlit.columns.Column): The Streamlit column to display the session state.
 
         Returns:
             None
         """
-  
+
         col.write("Current Model:")
         data = [{'Key': key, 'Value': str(value)} for key, value in st.session_state.items()]
         df = pd.DataFrame(data).reset_index(drop=True)