Update my_model/state_manager.py

my_model/state_manager.py

@@ -2,6 +2,7 @@ import pandas as pd
 import copy
 import time
 from PIL import Image
+from typing import Tuple, Dict
 import streamlit as st
 from my_model.utilities.gen_utilities import free_gpu_resources
 from my_model.KBVQA import KBVQA, prepare_kbvqa_model
@@ -10,20 +11,35 @@ from my_model.KBVQA import KBVQA, prepare_kbvqa_model
 
 
 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):
+        """
+        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])  
 
     def initialize_state(self):
+        """
+        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 'method' not in st.session_state:
-        #    st.session_state['method'] = "Fine-Tuned Model"
-        #if 'confidence_level' not in st.session_state:
-        #    st.session_state['confidence_level'] = 0.4
-        #if 'detection_model' not in st.session_state:
-        #    st.session_state['detection_model'] = 'detic'
         if 'images_data' not in st.session_state:
             st.session_state['images_data'] = {}
         if 'kbvqa' not in st.session_state:
@@ -38,18 +54,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'] = None
-
         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):
+    def set_up_widgets(self) -> None:
         """
         Sets up user interface widgets for selecting models, settings, and displaying model settings conditionally.
         """
@@ -68,7 +78,7 @@ class StateManager:
 
 
         
-    def set_slider_value(self, text, min_value, max_value, value, step, slider_key_name, col=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.
 
@@ -87,13 +97,18 @@ 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']
 
     def disable_widgets(self):
+        """
+        Disables widgets by setting the 'loading_in_progress' state to True.
+        """
+        
         st.session_state['loading_in_progress'] = True
+
         
     @property
     def settings_changed(self):
@@ -104,38 +119,30 @@ class StateManager:
             bool: True if any setting has changed, False otherwise.
         """
         return self.has_state_changed()
+
         
     @property
     def confidance_change(self):
-        return st.session_state["confidence_level"] != st.session_state["previous_state"]["confidence_level"]
-
-    @property
-    def display_model_settings(self):
         """
-        Displays a table of current model settings in the third column.
-    
+        Checks if the confidence level setting has changed compared to the previous state.
+
+        Returns:
+            bool: True if the confidence level has changed, False otherwise.
         """
-        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).reset_index(drop=True)
-        return self.col3.write(df)
+        
+        return st.session_state["confidence_level"] != st.session_state["previous_state"]["confidence_level"]
 
-    
-    def display_session_state(self):
+        
+    def update_prev_state(self):
         """
-        Displays a table of the complete application state..
+        Updates the 'previous_state' in the session state with the current state values.
         """
-  
-        st.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)
-        st.write(df)
         
-    def update_prev_state(self):
         for key in st.session_state['previous_state']:
             st.session_state['previous_state'][key] = st.session_state[key]
+
             
-    def load_model(self):
+    def load_model(self) -> None:
         """
         Loads the KBVQA model based on the chosen method and settings.
     
@@ -159,7 +166,17 @@ class StateManager:
         except Exception as e:
             st.error(f"Error loading model: {e}")
 
-    def force_reload_model(self):
+    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`.
+
+        - Deletes the current KBVQA model from the session state.
+        - 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.
+        """
+        
+        
         try:
             self.delete_model()
             free_gpu_resources()
@@ -174,9 +191,9 @@ class StateManager:
             st.error(f"Error reloading model: {e}")
             free_gpu_resources()
         
-    def delete_model(self):
+    def delete_model(self) -> None:
         """
-        Forces a reload of all models, freeing up GPU resources. This method deletes the current models and calls `free_gpu_resources`.
+        This method deletes the current models and calls `free_gpu_resources`.
         """
         
         free_gpu_resources()
@@ -191,7 +208,7 @@ class StateManager:
             
         
     # Function to check if any session state values have changed
-    def has_state_changed(self):
+    def has_state_changed(self) -> bool:
         """
         Compares current session state with the previous state to identify changes.
     
@@ -207,7 +224,7 @@ class StateManager:
         else: return False  # No changes found   
 
     
-    def get_model(self):
+    def get_model(self) -> KBVQA:
         """
         Retrieve the KBVQA model from the session state.
            
@@ -216,7 +233,7 @@ class StateManager:
         return st.session_state.get('kbvqa', None)
 
     @property
-    def is_model_loaded(self):
+    def is_model_loaded(self) -> bool:
         """
         Checks if the KBVQA model is loaded in the session state.
     
@@ -226,7 +243,7 @@ class StateManager:
         return 'kbvqa' in st.session_state and st.session_state['kbvqa'] is not None and st.session_state.kbvqa.all_models_loaded
 
     
-    def reload_detection_model(self):
+    def reload_detection_model(self) -> None:
         """
         Reloads only the detection model of the KBVQA model with updated settings.
     
@@ -251,7 +268,7 @@ class StateManager:
             st.error(f"Error reloading detection model: {e}")
 
     
-    def process_new_image(self, image_key, image):
+    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.
     
@@ -278,7 +295,7 @@ class StateManager:
             }
 
     
-    def analyze_image(self, image):
+    def analyze_image(self, image) -> Tuple[str, str, object]:
         """
         Analyzes the image using the KBVQA model.
 
@@ -302,7 +319,7 @@ class StateManager:
         return caption, detected_objects_str, image_with_boxes
 
     
-    def add_to_qa_history(self, image_key, question, answer, prompt_length):
+    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.
     
@@ -395,4 +412,27 @@ class StateManager:
             st.success(messae)
         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):
+        """
+        Displays a table of current model settings in the third column.
+    
+        """
+        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).reset_index(drop=True)
+        return self.col3.write(df)
+
+    
+    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).reset_index(drop=True)
+        st.write(df)