Update my_model/state_manager.py

my_model/state_manager.py

@@ -1,3 +1,7 @@
+# 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.
+
+
 import pandas as pd
 import copy
 import time
@@ -8,25 +12,25 @@ from my_model.utilities.gen_utilities import free_gpu_resources
 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.
 
+    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.
+
+    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.
+        col2 (streamlit.columns): The second column in the Streamlit layout.
+        col3 (streamlit.columns): The third column in the Streamlit layout.
+    """
 
-class StateManager:
     
-    # Hints for methods
-    # initialize_state: Initializes default values for session state.
-    # set_up_widgets: Creates UI elements for model selection and settings.
-    # set_slider_value: Generates a slider widget for numerical input.
-    # is_widget_disabled: Returns True if UI elements should be disabled.
-    # disable_widgets: Disables interactive UI elements during processing.
-    # settings_changed: Checks if any model settings have changed.
-    # confidance_change: Determines if the confidence level setting has changed.
-    # display_model_settings: Shows current model settings in the UI.
-    # display_session_state: Displays the current state of the application.
-    # update_prev_state: Updates the record of the previous application state.
-    # force_reload_model: Reloads the model, clearing and resetting necessary states.
-
-
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Initializes the StateManager instance, setting up the Streamlit columns for the user interface.
         """
@@ -34,10 +38,12 @@ class StateManager:
         # Create three columns with different widths
         self.col1, self.col2, self.col3 = st.columns([0.2, 0.6, 0.2])  
 
-    def initialize_state(self):
+    
+    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:
@@ -59,9 +65,13 @@ 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.
+
+        Returns:
+            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)
@@ -70,14 +80,11 @@ class StateManager:
         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)
         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:
         """
         Creates a slider widget with the specified parameters, optionally placing it in a specific column.
@@ -90,6 +97,9 @@ class StateManager:
             step (float): Step size for the slider.
             slider_key_name (str): Unique key for the slider.
             col (streamlit.columns.Column, optional): Column to place the slider in. Defaults to None (displayed in main area).
+
+        Returns:
+            None
         """
         
         if col is None:
@@ -99,30 +109,41 @@ class StateManager:
 
     
     @property
-    def is_widget_disabled(self):
+    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):
+    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):
+    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):
+    def confidance_change(self) -> bool:
         """
         Checks if the confidence level setting has changed compared to the previous state.
 
@@ -133,9 +154,12 @@ class StateManager:
         return st.session_state["confidence_level"] != st.session_state["previous_state"]["confidence_level"]
 
         
-    def update_prev_state(self):
+    def update_prev_state(self) -> None:
         """
         Updates the 'previous_state' in the session state with the current state values.
+
+        Returns:
+            None
         """
         
         for key in st.session_state['previous_state']:
@@ -151,6 +175,9 @@ class StateManager:
         - Sets the detection confidence level on the model object.
         - Updates previous state with current settings for change detection.
         - Updates the button label to "Reload Model".
+
+        Returns:
+            None
         """
         
         try:
@@ -166,6 +193,7 @@ class StateManager:
             
         except Exception as e:
             st.error(f"Error loading model: {e}")
+            
 
     def force_reload_model(self) -> None:
         """
@@ -175,9 +203,11 @@ class StateManager:
         - Calls `prepare_kbvqa_model` with `force_reload=True` to reload the model.
         - Updates the detection confidence level on the model object.
         - Displays a success message if the model is reloaded successfully.
+
+        Returns:
+            None
         """
         
-        
         try:
             self.delete_model()
             free_gpu_resources()
@@ -191,10 +221,14 @@ class StateManager:
         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`.
+
+        Returns:
+            None
         """
         
         free_gpu_resources()
@@ -210,11 +244,10 @@ class StateManager:
                 pass
             
         
-    # Function to check if any session state values have changed
     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.
         """
@@ -227,14 +260,17 @@ class StateManager:
         else: return False  # No changes found   
 
     
-    def get_model(self) -> KBVQA:
+    def get_model(self) -> KBVQA.KBVQA():
         """
-        Retrieve the KBVQA model from the session state.
-           
-        Returns: KBVQA object: The loaded KBVQA model, or None if not loaded.
+        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:
         """
@@ -243,6 +279,7 @@ class StateManager:
         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 
@@ -258,6 +295,9 @@ class StateManager:
         - Calls `prepare_kbvqa_model` with `only_reload_detection_model=True`.
         - Updates detection confidence level on the model object.
         - Displays a success message if model is reloaded successfully.
+
+        Returns:
+            None
         """
         
         try:
@@ -288,7 +328,9 @@ class StateManager:
         Args:
             image_key (str): Unique key for the image.
             image (obj): The uploaded image data.
-   
+
+        Returns:
+            None
         """
         
         if image_key not in st.session_state['images_data']:
@@ -300,6 +342,7 @@ class StateManager:
                 'analysis_done': False
             }
 
+
     
     def analyze_image(self, image) -> Tuple[str, str, object]:
         """
@@ -313,7 +356,6 @@ class StateManager:
         Args:
             image (obj): The image data to analyze.
 
-    
         Returns:
             tuple: A tuple containing the generated caption, detected objects string, and image with bounding boxes.
         """
@@ -329,28 +371,34 @@ class StateManager:
     
     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 hitory tracker.
+        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.
             answer (str): The answer generated by the KBVQA model.
+            prompt_length (int): The length of the prompt used for generating the answer.
+
+        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):
+    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, caption, detected_objects_str, analysis_done):
+    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.
     
@@ -359,6 +407,9 @@ class StateManager:
             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.
+
+        Returns:
+            None
         """
         if image_key in st.session_state['images_data']:
             st.session_state['images_data'][image_key].update({
@@ -368,18 +419,18 @@ class StateManager:
             })
 
  
-    def resize_image(self, image_input, new_width=None, new_height=None):
+    def resize_image(self, image_input, new_width: Optional[int] = None, new_height: Optional[int] = None) -> Image:
         """
-        Resize an image. If only new_width is provided, the height is adjusted to maintain aspect ratio.
+        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 (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.
+            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.
+            PIL.Image.Image: The resized image.
         """
     
         img = copy.deepcopy(image_input)
@@ -411,7 +462,18 @@ class StateManager:
 
     
  
-    def display_message(self, message, message_type):
+    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').
+
+        Returns:
+            None
+        """
+        
         if message_type == "warning":
             st.warning(message)
         elif message_type == "text":
@@ -424,10 +486,12 @@ class StateManager:
         
         
     @property
-    def display_model_settings(self):
+    def display_model_settings(self) -> None:
         """
         Displays a table of current model settings in the third column.
-    
+
+        Returns:
+            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' ]]
@@ -435,9 +499,15 @@ class StateManager:
         return self.col3.write(df)
 
     
-    def display_session_state(self, col):
+    def display_session_state(self, col) -> None:
         """
-        Displays a table of the complete application state..
+        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:")