Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -12,6 +12,7 @@ from my_model.object_detection import detect_and_draw_objects
 from my_model.captioner.image_captioning import get_caption
 from my_model.utilities.gen_utilities import free_gpu_resources
 from my_model.state_manager import StateManager
+from my_model.config import inference_config as config
 
 
 class InferenceRunner(StateManager):
@@ -19,8 +20,8 @@ class InferenceRunner(StateManager):
         
         super().__init__()
         self.initialize_state()
-        self.sample_images = [
-            "Files/sample1.jpg", "Files/sample2.jpg", "Files/sample3.jpg"]
+        self.sample_images = config.SAMPLE_IMAGES
+
 
     def answer_question(self, caption, detected_objects_str, question, model):
         free_gpu_resources()
@@ -76,16 +77,32 @@ class InferenceRunner(StateManager):
         
                     if image_data['analysis_done']:
                         st.session_state['loading_in_progress'] = False
-                        question = nested_col22.text_input(f"Ask a question about this image ({image_key[-11:]}):", key=f'question_{image_key}')
-                        if nested_col22.button('Get Answer', key=f'answer_{image_key}', on_click=self.disable_widgets, disabled=self.is_widget_disabled):
-
-                            if question not in [q for q, _ in qa_history]:
-                                answer = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
-                                st.session_state['loading_in_progress'] = False
-                                self.add_to_qa_history(image_key, question, answer)
-                            else: nested_col22.warning("This questions has already been answered.")
-                                
-                            st.session_state['loading_in_progress'] = False
+                        sample_questions = config.SAMPLE_QUESTIONS.get(image_key, [])
+                        selected_question = nested_col22.selectbox(
+                            "Select a sample question or type your own:",
+                            ["Custom question..."] + sample_questions,
+                            key=f'sample_question_{image_key}')
+                        
+                        # Text input for custom question
+                        custom_question = nested_col22.text_input(
+                            "Or ask your own question:",
+                            key=f'custom_question_{image_key}')
+                        # Use the selected sample question or the custom question
+                        question = custom_question if selected_question == "Custom question..." else selected_question
+                        
+
+                        if not question:
+                            nested_col22.warning("Please select or enter a question.")
+                        else:
+                            if question in [q for q, _ in qa_history]:
+                                nested_col22.warning("This question has already been answered.")
+                            else:
+                                if nested_col22.button('Get Answer', key=f'answer_{image_key}', disabled=self.is_widget_disabled):
+
+                                    answer = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
+                                    st.session_state['loading_in_progress'] = False
+                                    self.add_to_qa_history(image_key, question, answer)
+        
         
                     # Display Q&A history for each image
                     for num, (q, a) in enumerate(qa_history):
@@ -95,16 +112,9 @@ class InferenceRunner(StateManager):
         pass
 
 
-
-    
-
-
-
-
         
     def run_inference(self):
-        if st.session_state['loading_in_progress']:
-            st.rerun()
+        
         self.set_up_widgets()
 
         load_fine_tuned_model = False