Update my_model/utilities/st_utils.py

my_model/utilities/st_utils.py

@@ -92,8 +92,8 @@ class StateManager:
 
     def check_settings_changed(self, current_selected_method, current_detection_model, current_confidence_level):
         return (st.session_state['model_settings']['detection_model'] != current_detection_model or
-                st.session_state['model_settings']['confidence_level'] != current_confidence_level
-                st.session_state['model_settings']['selected_method'] != current_selected_method)
+                st.session_state['model_settings']['confidence_level'] != current_confidence_level or
+                st.session_state['selected_method'] != current_selected_method)
 
     def display_model_settings(self):
         st.write("### Current Model Settings:")
@@ -106,12 +106,9 @@ class StateManager:
         st.table(df)
 
     def is_model_loaded(self):
-        """Check if the model is loaded in the session state."""
         return 'kbvqa' in st.session_state and st.session_state['kbvqa'] is not None
 
-
     def reload_detection_model(self, detection_model, confidence_level):
-        """Reload only the detection model with new settings."""
         try:
             free_gpu_resources()
             if self.is_model_loaded():
@@ -122,12 +119,34 @@ class StateManager:
         except Exception as e:
             st.error(f"Error reloading detection model: {e}")
 
-
-
-
-
-
-
-
-
-
+    # New methods to be added
+    def process_new_image(self, image_key, image, kbvqa):
+        if image_key not in st.session_state['images_data']:
+            st.session_state['images_data'][image_key] = {
+                'image': image,
+                'caption': '',
+                'detected_objects_str': '',
+                'qa_history': [],
+                'analysis_done': False
+            }
+
+    def analyze_image(self, image, kbvqa):
+        img = copy.deepcopy(image)
+        caption = kbvqa.get_caption(img)
+        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({
+                'caption': caption,
+                'detected_objects_str': detected_objects_str,
+                'analysis_done': analysis_done
+            })