Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -16,14 +16,36 @@ from my_model.config import inference_config as config
 
 
 class InferenceRunner(StateManager):
+
+    """
+    InferenceRunner manages the user interface and interactions for a Streamlit-based
+    Knowledge-Based Visual Question Answering (KBVQA) application. It handles image uploads,
+    displays sample images, and facilitates the question-answering process using the KBVQA model.
+    it inherits the StateManager class.
+    """
+    
     def __init__(self):
+        """
+        Initializes the InferenceRunner instance, setting up the necessary state.
+        """
         
         super().__init__()
         self.initialize_state()
-        self.sample_images = config.SAMPLE_IMAGES
 
 
     def answer_question(self, caption, detected_objects_str, question, model):
+        """
+        Generates an answer to a given question based on the image's caption and detected objects.
+
+        Args:
+            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)
         free_gpu_resources()
@@ -31,10 +53,18 @@ class InferenceRunner(StateManager):
 
 
     def image_qa_app(self, kbvqa):
+        """
+        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(self.sample_images))
-        for idx, sample_image_path in enumerate(self.sample_images):
+        cols = self.col1.columns(len(config.SAMPLE_IMAGES))
+        for idx, sample_image_path in enumerate(config.SAMPLE_IMAGES):
             with cols[idx]:
                 image = Image.open(sample_image_path)
                 image_for_display = self.resize_image(sample_image_path, 80, 80)
@@ -42,13 +72,8 @@ class InferenceRunner(StateManager):
                 if st.button(f'Select Sample Image {idx + 1}', key=f'sample_{idx}'):
                     self.process_new_image(sample_image_path, image, kbvqa)
 
-                    
-
-        
         # 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)
 
@@ -67,7 +92,6 @@ class InferenceRunner(StateManager):
                                 caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image'], kbvqa)
                                 self.update_image_data(image_key, caption, detected_objects_str, True)
                             st.session_state['loading_in_progress'] = False
-
         
                     # Initialize qa_history for each image
                     qa_history = image_data.get('qa_history', [])
@@ -87,7 +111,6 @@ class InferenceRunner(StateManager):
                         # 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:
@@ -100,20 +123,19 @@ class InferenceRunner(StateManager):
                                     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):
                         nested_col22.text(f"Q{num+1}: {q}\nA{num+1}: {a}\n")
 
-    def display_message(self, message, warning=False, write=False, text=False):
-        pass
-
-
         
     def run_inference(self):
+        """
+        Sets up the widgets and manages the inference process. This method handles model loading,
+        reloading, and the overall flow of the inference process based on user interactions.
+
+        """
         
         self.set_up_widgets()
-
         load_fine_tuned_model = False
         fine_tuned_model_already_loaded = False
         reload_detection_model = False
@@ -123,27 +145,21 @@ class InferenceRunner(StateManager):
         if st.session_state['settings_changed']:
             self.col1.warning("Model settings have changed, please reload the model, this will take a second .. ")
            
-
         st.session_state.button_label = "Reload Model" if self.is_model_loaded() and self.settings_changed else "Load Model"
 
         with self.col1:
-                
             if st.session_state.method == "Fine-Tuned Model":
-                
                 with st.container():
                     nested_col11, nested_col12 = st.columns([0.5, 0.5])
                     if nested_col11.button(st.session_state.button_label, on_click=self.disable_widgets, disabled=self.is_widget_disabled):
-                        
                         if st.session_state.button_label == "Load Model":
                             if self.is_model_loaded():
                                 free_gpu_resources()
                                 fine_tuned_model_already_loaded = True
-                      
                             else:
                                 load_fine_tuned_model = True
                         else:
                             reload_detection_model = True
-                            
                     if nested_col12.button("Force Reload", on_click=self.disable_widgets, disabled=self.is_widget_disabled):
                         force_reload_full_model = True
                         
@@ -172,14 +188,12 @@ class InferenceRunner(StateManager):
                     st.session_state['time_taken_to_load_model'] = int(time.time()-t1)
                     st.session_state['loading_in_progress'] = False
                     st.session_state['model_loaded'] = True
-                
+                    
             elif st.session_state.method == "In-Context Learning (n-shots)":
                 self.col1.warning(f'Model using {st.session_state.method} is not deployed yet, will be ready later.')
                 st.session_state['loading_in_progress'] = False
-
         
         if self.is_model_loaded():
-
             free_gpu_resources()
             st.session_state['loading_in_progress'] = False
             self.image_qa_app(self.get_model())