Update my_model/state_manager.py

my_model/state_manager.py

@@ -29,6 +29,9 @@ class StateManager:
 
 
     def set_up_widgets(self):
+        """
+        Sets up user interface widgets for selecting models, settings, and displaying model settings conditionally.
+        """
 
         self.col1.selectbox("Choose a method:", ["Fine-Tuned Model", "In-Context Learning (n-shots)"], index=0, key='method')
         detection_model = self.col1.selectbox("Choose a model for objects detection:", ["yolov5", "detic"], index=1, key='detection_model')
@@ -45,6 +48,19 @@ class StateManager:
 
         
     def set_slider_value(self, text, min_value, max_value, value, step, slider_key_name, col=None):
+        """
+        Creates a slider widget with the specified parameters, optionally placing it in a specific column.
+
+        Args:
+            text (str): Text to display next to the slider.
+            min_value (float): Minimum value for the slider.
+            max_value (float): Maximum value for the slider.
+            value (float): Initial value for the slider.
+            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).
+        """
+        
         if col is None:
             return st.slider(text, min_value, max_value, value, step, key=slider_key_name)
         else:
@@ -53,10 +69,21 @@ class StateManager:
         
     @property
     def settings_changed(self):
+        """
+        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()
 
     
     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 = [{'Key': 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', ]]
         df = pd.DataFrame(data)
@@ -65,6 +92,10 @@ class StateManager:
 
     
     def display_session_state(self):
+        """
+        Displays a table of the complete application state..
+        """
+  
         st.write("Current Model:")
         data = [{'Key': key, 'Value': str(value)} for key, value in st.session_state.items()]
         df = pd.DataFrame(data)
@@ -72,7 +103,16 @@ class StateManager:
         
 
     def load_model(self):
-        """Load the KBVQA model with specified settings."""
+        """
+        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.
+        - Updates previous state with current settings for change detection.
+        - Updates the button label to "Reload Model".
+        """
+        
         try:
             free_gpu_resources()
             st.session_state['kbvqa'] = prepare_kbvqa_model()
@@ -91,6 +131,12 @@ class StateManager:
             
     # Function to check if any session state values have changed
     def has_state_changed(self):
+        """
+        Compares current session state with the previous state to identify changes.
+    
+        Returns:
+            bool: True if any change is found, False otherwise.
+        """
         for key in st.session_state['previous_state']:
             if st.session_state[key] != st.session_state['previous_state'][key]:
                 return True  # Found a change
@@ -98,15 +144,35 @@ class StateManager:
 
     
     def get_model(self):
-        """Retrieve the KBVQA model from the session state."""
+        """
+        Retrieve the KBVQA model from the session state.
+           
+        Returns: KBVQA object: The loaded KBVQA model, or None if not loaded.
+        """
         return st.session_state.get('kbvqa', None)
 
     
     def is_model_loaded(self):
+        """
+        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
 
     
     def reload_detection_model(self):
+        """
+        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`.
+        - Updates detection confidence level on the model object.
+        - Displays a success message if model is reloaded successfully.
+        """
+        
         try:
             free_gpu_resources()
             if self.is_model_loaded():
@@ -119,6 +185,22 @@ class StateManager:
 
     
     def process_new_image(self, image_key, image, kbvqa):
+        """
+        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.
+            kbvqa (KBVQA object): The loaded KBVQA model.
+        """
+        
         if image_key not in st.session_state['images_data']:
             st.session_state['images_data'][image_key] = {
                 'image': image,
@@ -130,6 +212,21 @@ class StateManager:
 
     
     def analyze_image(self, image, kbvqa):
+        """
+        Analyzes the image using the KBVQA model.
+
+        - Creates a copy of the image to avoid modifying the original.
+        - 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.
+            kbvqa (KBVQA object): The loaded KBVQA model.
+    
+        Returns:
+            tuple: A tuple containing the generated caption, detected objects string, and image with bounding boxes.
+        """
         img = copy.deepcopy(image)
         st.text("Analyzing the image .. ")
         caption = kbvqa.get_caption(img)
@@ -138,22 +235,60 @@ class StateManager:
 
     
     def add_to_qa_history(self, image_key, question, answer):
+        """
+        Adds a question-answer pair to the QA history of a specific image, to be used as hitory 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.
+        """
         if image_key in st.session_state['images_data']:
             st.session_state['images_data'][image_key]['qa_history'].append((question, answer))
 
     
     def get_images_data(self):
+        """
+        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']
 
     @staticmethod
     def resize_image(image_path, new_width, new_height):
-        """Resize an image to the specified width and height."""
-        image = Image.open(image_path)
-        resized_image = image.resize((new_width, new_height))
+        """
+        Resizes an image from the specified to the given dimensions.
+    
+        Args:
+            image_path (str): Path to the image file.
+            new_width (int): Desired width for the resized image.
+            new_height (int): Desired height for the resized image.
+    
+        Returns:
+            Image: The resized image object.
+        """
+
+        if isinstance(image_path, str):
+            # Open the image from a file path
+            image = Image.open(image_path)
+        elif hasattr(image_path, 'read'):
+            resized_image = image.resize((new_width, new_height))
+            
         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,