Update my_model/state_manager.py

my_model/state_manager.py

@@ -4,12 +4,15 @@ import streamlit as st
 from my_model.utilities.gen_utilities import free_gpu_resources
 from my_model.KBVQA import KBVQA, prepare_kbvqa_model
 
-# Create two columns with different widths
-col1, col2, col3 = st.columns([0.2, 0.6, 0.2])  
+
 
 
 class StateManager:
 
+    def __init__(self):
+        # Create three columns with different widths
+        self.col1, self.col2, self.col3 = st.columns([0.2, 0.6, 0.2])  
+
     def initialize_state(self):
         if 'images_data' not in st.session_state:
             st.session_state['images_data'] = {}
@@ -26,18 +29,18 @@ class StateManager:
 
     def set_up_widgets(self):
 
-        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')
-            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.1, slider_key_name='confidence_level')
+    
+        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')
+        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.1, slider_key_name='confidence_level', col=self.col1)
 
         # Conditional display of model settings
 
-        with col3:
-            show_model_settings = st.checkbox("Show Model Settings", False)
-            if show_model_settings:
-                self.display_model_settings()
+ 
+        show_model_settings = self.col3.checkbox("Show Model Settings", False)
+        if show_model_settings:
+            self.display_model_settings()
 
 
         
@@ -54,11 +57,11 @@ class StateManager:
 
     
     def display_model_settings(self):
-        st.write("##### Current Model Settings:")
+        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)
         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')]}])
-        col3.table(styled_df)
+        self.col3.table(styled_df)
 
     
     def display_session_state(self):
@@ -109,7 +112,7 @@ class StateManager:
             if self.is_model_loaded():
                 prepare_kbvqa_model(only_reload_detection_model=True)
                 st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level
-                col1.success("Model reloaded with updated settings and ready for inference.")
+                self.col1.success("Model reloaded with updated settings and ready for inference.")
             free_gpu_resources()
         except Exception as e:
             st.error(f"Error reloading detection model: {e}")