Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -33,7 +33,7 @@ class InferenceRunner(StateManager):
        # self.initialize_state()
 
 
-    def answer_question(self, caption, detected_objects_str, question, model):
+    def answer_question(self, caption, detected_objects_str, question):
         """
         Generates an answer to a given question based on the image's caption and detected objects.
 
@@ -41,27 +41,25 @@ class InferenceRunner(StateManager):
             caption (str): The caption generated for the image.
             detected_objects_str (str): String representation of objects detected in the image.
             question (str): The user's question about the image.
-            model (KBVQA): The loaded KBVQA model used for generating the answer.
+  
 
         Returns:
             str: The generated answer to the question.
         """
         free_gpu_resources()
-        answer = model.generate_answer(question, caption, detected_objects_str)
+        answer = self.session_state.kbvqa.generate_answer(question, caption, detected_objects_str)
         prompt_length  = model.current_prompt_length
         free_gpu_resources()
         return answer, prompt_length
 
 
-    def image_qa_app(self, kbvqa):
+    def image_qa_app(self):
         """
         Main application interface for image-based question answering. It handles displaying
         of sample images, uploading of new images, and facilitates the QA process.
-
-        Args:
-            kbvqa (KBVQA): The loaded KBVQA model used for image analysis and question answering.
         """
         
+        
         # Display sample images as clickable thumbnails
         self.col1.write("Choose from sample images:")
         cols = self.col1.columns(len(config.SAMPLE_IMAGES))
@@ -71,12 +69,12 @@ class InferenceRunner(StateManager):
                 image_for_display = self.resize_image(sample_image_path, 80, 80)
                 st.image(image_for_display)
                 if st.button(f'Select Sample Image {idx + 1}', key=f'sample_{idx}'):
-                    self.process_new_image(sample_image_path, image, kbvqa)
+                    self.process_new_image(sample_image_path, image)
 
         # Image uploader
         uploaded_image = self.col1.file_uploader("Or upload an Image", type=["png", "jpg", "jpeg"])
         if uploaded_image is not None:
-            self.process_new_image(uploaded_image.name, Image.open(uploaded_image), kbvqa)
+            self.process_new_image(uploaded_image.name, Image.open(uploaded_image))
 
         # Display and interact with each uploaded/selected image
         self.display_session_state()
@@ -91,7 +89,7 @@ class InferenceRunner(StateManager):
                         with nested_col22:
                             if st.button('Analyze Image', key=f'analyze_{image_key}', on_click=self.disable_widgets, disabled=self.is_widget_disabled):
                                 
-                                caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image'], kbvqa)
+                                caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image'])
                                 self.update_image_data(image_key, caption, detected_objects_str, True)
                             st.session_state['loading_in_progress'] = False
         
@@ -121,7 +119,7 @@ class InferenceRunner(StateManager):
                         else:
                             if nested_col22.button('Get Answer', key=f'answer_{image_key}', disabled=self.is_widget_disabled):
 
-                                answer, prompt_length = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
+                                answer, prompt_length = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question)
                                 st.session_state['loading_in_progress'] = False
                                 self.add_to_qa_history(image_key, question, answer, prompt_length)
         
@@ -198,6 +196,6 @@ class InferenceRunner(StateManager):
         if self.is_model_loaded:
             free_gpu_resources()
             st.session_state['loading_in_progress'] = False
-            self.image_qa_app(self.get_model())
+            self.image_qa_app()