Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -51,7 +51,6 @@ class InferenceRunner(StateManager):
 
 
     def display_sample_images(self):
-        # Display sample images as clickable thumbnails
         self.col1.write("Choose from sample images:")
         cols = self.col1.columns(len(config.SAMPLE_IMAGES))
         for idx, sample_image_path in enumerate(config.SAMPLE_IMAGES):
@@ -62,83 +61,54 @@ class InferenceRunner(StateManager):
                 if st.button(f'Select Sample Image {idx + 1}', key=f'sample_{idx}'):
                     self.process_new_image(sample_image_path, image)
 
-    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.
-        """
-        
-        self.display_sample_images()
-        # Image uploader
+    def handle_image_upload(self):
         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))
 
-        # Display and interact with each uploaded/selected image
-        
+    def display_image_and_analysis(self, image_key, image_data):
+        nested_col21, nested_col22 = st.columns([0.65, 0.35])
+        image_for_display = self.resize_image(image_data['image'], 600)
+        nested_col21.image(image_for_display, caption=f'Uploaded Image: {image_key[-11:]}')
+        self.handle_analysis_button(image_key, image_data, nested_col22)
+
+    def handle_analysis_button(self, image_key, image_data, nested_col22):
+        if not image_data['analysis_done'] or self.settings_changed or self.confidance_change:
+            nested_col22.text("Please click 'Analyze Image'..")
+            analyze_button_key = f'analyze_{image_key}_{st.session_state.detection_model}_{st.session_state.confidence_level}'
+            if st.button('Analyze Image', key=analyze_button_key, on_click=self.disable_widgets, disabled=self.is_widget_disabled):
+                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
+
+    def handle_question_answering(self, image_key, image_data, nested_col22):
+        if image_data['analysis_done']:
+            self.display_question_answering_interface(image_key, image_data, nested_col22)
+
+    def display_question_answering_interface(self, image_key, image_data, nested_col22):
+        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}')
+        custom_question = nested_col22.text_input("Or ask your own question:", key=f'custom_question_{image_key}')
+        question = custom_question if selected_question == "Custom question..." else selected_question
+        self.process_question(image_key, question, image_data, nested_col22)
+
+    def process_question(self, image_key, question, image_data, nested_col22):
+        qa_history = image_data.get('qa_history', [])
+        if question and (question not in [q for q, _, _ in qa_history] or self.settings_changed or self.confidance_change):
+            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)
+                self.add_to_qa_history(image_key, question, answer, prompt_length)
+                nested_col22.text(f"Q: {question}\nA: {answer}\nPrompt Length: {prompt_length}")
+
+    def image_qa_app(self):
+        self.display_sample_images()
+        self.handle_image_upload()
+        self.display_session_state()
         with self.col2:
             for image_key, image_data in self.get_images_data().items():
-                
                 with st.container():
-                    nested_col21, nested_col22 = st.columns([0.65, 0.35])
-                    image_for_display = self.resize_image(image_data['image'], 600)
-                    nested_col21.image(image_for_display, caption=f'Uploaded Image: {image_key[-11:]}')
-                    
-                    if not image_data['analysis_done'] or self.settings_changed or self.confidance_change:  # if not done analysis before or even done but settings changed, then we need to analyze again
-                     
-                        nested_col22.text("Please click 'Analyze Image'..")
-                        free_gpu_resources()
-                        with nested_col22:
-                        
-                            analyze_button_key = f'analyze_{image_key}_{st.session_state.detection_model}_{st.session_state.confidence_level}'  # unique key for each click
-                            
-                            if st.button('Analyze Image', key=analyze_button_key, on_click=self.disable_widgets, disabled=self.is_widget_disabled):
-                                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
-                            free_gpu_resources()
-                            
-        
-                    # Initialize qa_history for each image
-                    qa_history = image_data.get('qa_history', [])
-                 
-                    if image_data['analysis_done']:
-                    
-                        free_gpu_resources()
-                        if self.confidance_change:
-                        
-                            nested_col22.warning("If you change the Confidence level, please click analyze again.")
-                            
-                        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 question in [q for q, _, _ in qa_history] and not self.settings_changed and not self.confidance_change:
-                            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):
-                                free_gpu_resources()
-                           
-                                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)
-                    
-                    # Display Q&A history and prompts lengths for each image
-                    for num, (q, a, p) in enumerate(qa_history):
-                        nested_col22.text(f"Q{num+1}: {q}\nA{num+1}: {a}\nPrompt Length: {p}\n")
-                        free_gpu_resources()
-                     
+                    self.display_image_and_analysis(image_key, image_data)
+                    self.handle_question_answering(image_key, image_data, nested_col22)