Update my_model/state_manager.py

my_model/state_manager.py

@@ -24,7 +24,7 @@ class StateManager:
     def set_up_widgets(self):
         
         # Create two columns with different widths
-        col1, col2 = st.columns([0.8, 0.2])  # Adjust the ratio as needed
+        col1, col2, col3 = st.columns([0.2, 0.6, 0.2])  # Adjust the ratio as needed
         with col1:
             st.selectbox("Choose a method:", ["Fine-Tuned Model", "In-Context Learning (n-shots)"], index=0, key='method')
             detection_model = st.selectbox("Choose a model for objects detection:", ["yolov5", "detic"], index=1, key='detection_model')
@@ -37,15 +37,19 @@ class StateManager:
             if show_model_settings:
                 self.display_model_settings()
 
+        col3.header("COL3")
+
         
     def set_slider_value(self, text, min_value, max_value, value, step, slider_key_name):
 
         return st.slider(text, min_value, max_value, value, step, key=slider_key_name)
+
         
     @property
     def settings_changed(self):
         return self.has_state_changed()
 
+    
     def display_model_settings(self):
         st.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', ]]
@@ -53,11 +57,13 @@ class StateManager:
         styled_df = df.style.set_properties(**{'background-color': 'black', 'color': 'white', 'border-color': 'white'}).set_table_styles([{'selector': 'th','props': [('background-color', 'black'), ('font-weight', 'bold')]}])
         st.table(styled_df)
 
+    
     def display_session_state(self):
         st.write("Current Model:")
         data = [{'Key': key, 'Value': str(value)} for key, value in st.session_state.items()]
         df = pd.DataFrame(data)
         st.table(df)
+        
 
     def load_model(self):
         """Load the KBVQA model with specified settings."""
@@ -76,6 +82,7 @@ class StateManager:
         except Exception as e:
             st.error(f"Error loading model: {e}")
 
+            
     # Function to check if any session state values have changed
     def has_state_changed(self):
         for key in st.session_state['previous_state']:
@@ -83,13 +90,16 @@ class StateManager:
                 return True  # Found a change
         else: return False  # No changes found   
 
+    
     def get_model(self):
         """Retrieve the KBVQA model from the session state."""
         return st.session_state.get('kbvqa', None)
 
+    
     def is_model_loaded(self):
         return 'kbvqa' in st.session_state and st.session_state['kbvqa'] is not None
 
+    
     def reload_detection_model(self):
         try:
             free_gpu_resources()
@@ -112,6 +122,7 @@ class StateManager:
                 'analysis_done': False
             }
 
+    
     def analyze_image(self, image, kbvqa):
         img = copy.deepcopy(image)
         st.text("Analyzing the image .. ")
@@ -119,13 +130,16 @@ class StateManager:
         image_with_boxes, detected_objects_str = kbvqa.detect_objects(img)
         return caption, detected_objects_str, image_with_boxes
 
+    
     def add_to_qa_history(self, image_key, question, answer):
         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):
         return st.session_state['images_data']
 
+    
     def update_image_data(self, image_key, caption, detected_objects_str, analysis_done):
         if image_key in st.session_state['images_data']:
             st.session_state['images_data'][image_key].update({